def start_oauth2_dance(domain):
current_user = users.get_current_user()
user_email = current_user.email()
login_hint = user_email
user_domain = current_user.email().split('@')[1]
if user_domain != domain:
domain = user_domain
primary_domain = PrimaryDomain.get_or_create(domain)
approval_prompt = 'auto' if primary_domain.refresh_token else 'force'
scope = constants.OAUTH2_SCOPES
state = urllib.quote(domain)
redirect_uri = url_for('oauth_callback', _external=True)
oauth_helper = OAuthDanceHelper(scope=scope, redirect_uri=redirect_uri,
approval_prompt=approval_prompt)
url = oauth_helper.step1_get_authorize_url()
# TODO: Add a random token to avoid forgery
return redirect("%s&state=%s&login_hint=%s" % (url, state, login_hint))
def schedule_user_move(user_email=None, tag=None, move_process_key=None,
domain_name=None):
if domain_name:
try:
primary_domain = PrimaryDomain.get_or_create(
domain_name=domain_name)
if primary_domain.credentials:
email_settings_helper = EmailSettingsHelper(
credentials_json=primary_domain.credentials,
domain=domain_name,
refresh_token=primary_domain.refresh_token
)
email_settings_helper.enable_imap(user_email)
logging.info('IMAP enabled for [%s]',
user_email)
else:
logging.warn('Error trying to enable IMAP for user [%s]',
user_email)
except:
logging.exception('Domain [%s] is not authorized, IMAP not enabled',
domain_name)
for chunk_ids in get_messages(user_email=user_email, tag=tag,
process_id=move_process_key.id()):
if len(chunk_ids) > 0:
new_chunk_ids = []
chunk_sizes = []
for chunk in chunkify(chunk_ids,
chunk_size=constants.MESSAGE_BATCH_SIZE):
new_chunk_ids.append([chunk[0], chunk[-1]])
chunk_sizes.append(len(chunk))
logging.info('Scheduling user [%s] messages move', user_email)
deferred.defer(move_messages, user_email=user_email, tag=tag,
chunk_ids=new_chunk_ids,
process_id=move_process_key.id(),
chunk_sizes=chunk_sizes)
def oauth_callback():
code = request.args.get('code', None)
if not code:
logging.error('No code, no authorization')
abort(500)
state = request.args.get('state', None)
if not state:
logging.error('No state, no authorization')
abort(500)
domain_name = urllib.unquote(state)
redirect_uri = url_for('oauth_callback', _external=True)
oauth_helper = OAuthDanceHelper(redirect_uri=redirect_uri)
credentials = oauth_helper.step2_exchange(code)
primary_domain = PrimaryDomain.get_or_create(domain_name)
primary_domain.credentials = credentials.to_json()
if credentials.refresh_token:
primary_domain.refresh_token = credentials.refresh_token
user = users.get_current_user()
primary_domain.admin_email = user.email()
primary_domain.put()
return redirect(url_for('settings'))
def clean_messages(user_email=None, password=None, chunk_ids=list(),
retry_count=0, process_id=None):
cleaned_successfully = []
remaining = []
if len(chunk_ids) <= 0:
process = CleanUserProcess.get_by_id(process_id)
process.status = constants.FINISHED
process.put()
return True
try:
process = CleanUserProcess.get_by_id(process_id)
imap = IMAPHelper()
imap.login(email=user_email, password=process.source_password)
imap.select()
domain_name = user_email.split('@')[1]
primary_domain = PrimaryDomain.get_or_create(
domain_name)
try:
drive = DriveHelper(credentials_json=primary_domain.credentials,
admin_email=primary_domain.admin_email,
refresh_token=primary_domain.refresh_token)
folder = drive.get_folder(constants.ATTACHMENT_FOLDER)
if not folder:
folder = drive.create_folder(constants.ATTACHMENT_FOLDER)
sub_folder = drive.get_folder(user_email)
if not sub_folder:
sub_folder = drive.create_folder(user_email,
[{'id': folder['id']}])
except Exception as e:
logging.error(
"Couldn't authenticate drive for user %s" % user_email)
raise e
try:
migration = MigrationHelper(
credentials_json=primary_domain.credentials,
refresh_token=primary_domain.refresh_token)
except Exception as e:
logging.error(
"Couldn't authenticate migration api for user %s" % user_email)
raise e
for message_id in chunk_ids:
try:
result = clean_message(msg_id=message_id, imap=imap,
drive=drive,
migration=migration,
folder_id=sub_folder['id'],
user_email=user_email,
process_id=process_id)
if result:
counter.load_and_increment_counter(
'cleaning_%s_ok_count' % (user_email),
namespace=str(process_id))
cleaned_successfully.append(message_id)
else:
counter.load_and_increment_counter(
'cleaning_%s_error_count' % user_email,
namespace=str(process_id))
logging.error(
'Error cleaning message ID [%s] for user [%s]: [%s] ',
message_id, user_email, result)
except Exception as e:
logging.exception(
'Failed cleaning individual message ID [%s] for user [%s]',
message_id, user_email)
remaining = []
if retry_count < constants.MAX_CLEAN_RETRIES:
for chunk_msg in chunk_ids:
if chunk_msg not in cleaned_successfully:
remaining.append(chunk_msg)
logging.info(
'Scheduling [%s] remaining cleaning messages for user [%s]',
len(remaining), user_email)
deferred.defer(clean_messages, user_email=user_email,
chunk_ids=remaining,
process_id=process_id,
retry_count=retry_count + 1)
else:
for chunk_msg in chunk_ids:
if message_id == chunk_msg:
continue
if chunk_msg not in cleaned_successfully:
remaining.append(chunk_msg)
logging.info(
'Giving up cleaning message [%s] for '
'user [%s]', message_id,
user_email)
counter.load_and_increment_counter(
'cleaning_%s_error_count' % user_email,
delta=1,
namespace=str(process_id))
deferred.defer(clean_messages, user_email=user_email,
chunk_ids=remaining,
process_id=process_id)
break
except Exception as e:
logging.exception('Failed cleaning messages chunk')
raise e
finally:
if imap:
imap.close()
if len(chunk_ids) < 10 or (len(cleaned_successfully) + 10 > len(chunk_ids)):
process.status = constants.FINISHED
process.put()
def list_process():
form = CleanUserProcessForm()
user = users.get_current_user()
clean_process_saved = False
clean_processes = []
clean_process_query = CleanUserProcess.query(CleanUserProcess.owner_email == user.email()).order()
query_params = {}
if request.method == 'POST':
if form.validate_on_submit():
primary_domain = PrimaryDomain.get_or_create(
domain_name = user.email().split('@')[1])
logged_in = 'NO'
current_user = users.get_current_user()
if current_user.email() == primary_domain.admin_email:
imap = IMAPHelper()
logged_in, _ = imap.login(
form.data['source_email'], form.data['source_password'])
imap.close()
if logged_in != 'OK':
form.source_email.errors.append(
"Can't access the email with those credentials")
else:
clean_user_process = CleanUserProcess(
owner_email=user.email(),
destination_message_email=user.email(),
status=constants.STARTED
)
for key, value in form.data.iteritems():
setattr(clean_user_process, key, value)
clean_process_key = clean_user_process.put()
clean_process_saved = True
# TODO: process does not appears immediately after it's saved
# launch Pipeline
deferred.defer(schedule_user_cleaning, user_email=form.data['source_email'],
process_id=clean_process_key.id())
is_prev = request.args.get('prev', False)
url_cursor = request.args.get('cursor', None)
cursor = Cursor(urlsafe=url_cursor) if url_cursor else None
if is_prev:
clean_process_query = clean_process_query.order(
CleanUserProcess.created)
cursor = cursor.reversed()
else:
clean_process_query = clean_process_query.order(
-CleanUserProcess.created)
data, next_curs, more = clean_process_query.fetch_page(
constants.PAGE_SIZE, start_cursor=cursor)
clean_processes.extend(data)
if is_prev:
prev_curs = next_curs.reversed().urlsafe() if more else None
next_curs = url_cursor
else:
prev_curs = url_cursor
next_curs = next_curs.urlsafe() if more else None
return render_template('process.html', form=form, user=user.email(),
clean_process_saved=clean_process_saved,
clean_processes=clean_processes, next_curs=next_curs,
more=more, prev_curs=prev_curs)
