From 9115ff89f3754e2cf952a8821c1f2469c9fca55b Mon Sep 17 00:00:00 2001 From: Micah Lee Date: Mon, 11 Feb 2019 22:46:39 -0800 Subject: Refactor receive mode to start saving files to data_dir with .part extension while they're downloading --- onionshare/web/receive_mode.py | 172 ++++++++++++++++++++--------------------- onionshare/web/web.py | 2 +- 2 files changed, 83 insertions(+), 91 deletions(-) diff --git a/onionshare/web/receive_mode.py b/onionshare/web/receive_mode.py index f035271a..be68900e 100644 --- a/onionshare/web/receive_mode.py +++ b/onionshare/web/receive_mode.py @@ -58,104 +58,37 @@ class ReceiveModeWeb(object): def upload_logic(slug_candidate=''): """ - Upload files. + Handle the upload files POST request, though at this point, the files have + already been uploaded and saved to their correct locations. """ - # Figure out what the receive mode dir should be - now = datetime.now() - date_dir = now.strftime("%Y-%m-%d") - time_dir = now.strftime("%H.%M.%S") - receive_mode_dir = os.path.join(self.common.settings.get('data_dir'), date_dir, time_dir) - valid = True - try: - os.makedirs(receive_mode_dir, 0o700, exist_ok=True) - except PermissionError: - self.web.add_request(self.web.REQUEST_ERROR_DATA_DIR_CANNOT_CREATE, request.path, { - "receive_mode_dir": receive_mode_dir - }) - print(strings._('error_cannot_create_data_dir').format(receive_mode_dir)) - valid = False - if not valid: - flash('Error uploading, please inform the OnionShare user', 'error') - if self.common.settings.get('public_mode'): - return redirect('/') - else: - return redirect('/{}'.format(slug_candidate)) - files = request.files.getlist('file[]') filenames = [] - print('') for f in files: if f.filename != '': - # Automatically rename the file, if a file of the same name already exists filename = secure_filename(f.filename) filenames.append(filename) - local_path = os.path.join(receive_mode_dir, filename) - if os.path.exists(local_path): - if '.' in filename: - # Add "-i", e.g. change "foo.txt" to "foo-2.txt" - parts = filename.split('.') - name = parts[:-1] - ext = parts[-1] - - i = 2 - valid = False - while not valid: - new_filename = '{}-{}.{}'.format('.'.join(name), i, ext) - local_path = os.path.join(receive_mode_dir, new_filename) - if os.path.exists(local_path): - i += 1 - else: - valid = True - else: - # If no extension, just add "-i", e.g. change "foo" to "foo-2" - i = 2 - valid = False - while not valid: - new_filename = '{}-{}'.format(filename, i) - local_path = os.path.join(receive_mode_dir, new_filename) - if os.path.exists(local_path): - i += 1 - else: - valid = True - + local_path = os.path.join(request.receive_mode_dir, filename) basename = os.path.basename(local_path) - if f.filename != basename: - # Tell the GUI that the file has changed names - self.web.add_request(self.web.REQUEST_UPLOAD_FILE_RENAMED, request.path, { - 'id': request.upload_id, - 'old_filename': f.filename, - 'new_filename': basename - }) # Tell the GUI the receive mode directory for this file self.web.add_request(self.web.REQUEST_UPLOAD_SET_DIR, request.path, { 'id': request.upload_id, 'filename': basename, - 'dir': receive_mode_dir + 'dir': request.receive_mode_dir }) - # Make sure receive mode dir exists before writing file - valid = True - try: - os.makedirs(receive_mode_dir, 0o700, exist_ok=True) - except PermissionError: - self.web.add_request(self.web.REQUEST_ERROR_DATA_DIR_CANNOT_CREATE, request.path, { - "receive_mode_dir": receive_mode_dir - }) - print(strings._('error_cannot_create_data_dir').format(receive_mode_dir)) - valid = False - if not valid: - flash('Error uploading, please inform the OnionShare user', 'error') - if self.common.settings.get('public_mode'): - return redirect('/') - else: - return redirect('/{}'.format(slug_candidate)) - self.common.log('ReceiveModeWeb', 'define_routes', '/upload, uploaded {}, saving to {}'.format(f.filename, local_path)) print(strings._('receive_mode_received_file').format(local_path)) - f.save(local_path) - # Note that flash strings are on English, and not translated, on purpose, + if request.upload_error: + self.common.log('ReceiveModeWeb', 'define_routes', '/upload, there was an upload error') + flash('Error uploading, please inform the OnionShare user', 'error') + if self.common.settings.get('public_mode'): + return redirect('/') + else: + return redirect('/{}'.format(slug_candidate)) + + # Note that flash strings are in English, and not translated, on purpose, # to avoid leaking the locale of the OnionShare user if len(filenames) == 0: flash('No files uploaded', 'info') @@ -198,7 +131,6 @@ class ReceiveModeWeb(object): return upload_logic() - class ReceiveModeWSGIMiddleware(object): """ Custom WSGI middleware in order to attach the Web object to environ, so @@ -214,10 +146,11 @@ class ReceiveModeWSGIMiddleware(object): return self.app(environ, start_response) -class ReceiveModeTemporaryFile(object): +class ReceiveModeFile(object): """ - A custom TemporaryFile that tells ReceiveModeRequest every time data gets - written to it, in order to track the progress of uploads. + A custom file object that tells ReceiveModeRequest every time data gets + written to it, in order to track the progress of uploads. It starts out with + a .part file extension, and when it's complete it removes that extension. """ def __init__(self, request, filename, write_func, close_func): self.onionshare_request = request @@ -225,8 +158,17 @@ class ReceiveModeTemporaryFile(object): self.onionshare_write_func = write_func self.onionshare_close_func = close_func - # Create a temporary file - self.f = tempfile.TemporaryFile('wb+') + self.filename = os.path.join(self.onionshare_request.receive_mode_dir, secure_filename(filename)) + self.filename_in_progress = '{}.part'.format(self.filename) + + # Open the file + try: + self.f = open(self.filename_in_progress, 'wb+') + except: + # This will only happen if someone is messaging with the data dir while + # OnionShare is running, but if it does make sure to throw an error + self.upload_error = True + self.f = tempfile.TemporaryFile('wb+') # Make all the file-like methods and attributes actually access the # TemporaryFile, except for write @@ -241,7 +183,7 @@ class ReceiveModeTemporaryFile(object): """ Custom write method that calls out to onionshare_write_func """ - if not self.onionshare_request.stop_q.empty(): + if self.upload_error or (not self.onionshare_request.stop_q.empty()): self.close() self.onionshare_request.close() return @@ -254,6 +196,11 @@ class ReceiveModeTemporaryFile(object): Custom close method that calls out to onionshare_close_func """ self.f.close() + + if not self.upload_error: + # Rename the in progress file to the final filename + os.rename(self.filename_in_progress, self.filename) + self.onionshare_close_func(self.onionshare_filename) @@ -283,6 +230,49 @@ class ReceiveModeRequest(Request): self.upload_request = True if self.upload_request: + # No errors yet + self.upload_error = False + + # Figure out what files should be saved + now = datetime.now() + date_dir = now.strftime("%Y-%m-%d") + time_dir = now.strftime("%H.%M.%S") + self.receive_mode_dir = os.path.join(self.web.common.settings.get('data_dir'), date_dir, time_dir) + + # Create that directory, which shouldn't exist yet + try: + os.makedirs(self.receive_mode_dir, 0o700, exist_ok=False) + except OSError: + # If this directory already exists, maybe someone else is uploading files at + # the same second, so use a different name in that case + if os.path.exists(self.receive_mode_dir): + # Keep going until we find a directory name that's available + i = 1 + while True: + new_receive_mode_dir = '{}-{}'.format(self.receive_mode_dir, i) + try: + os.makedirs(new_receive_mode_dir, 0o700, exist_ok=False) + break + except OSError: + pass + i += 1 + # Failsafe + if i == 100: + self.web.common.log('ReceiveModeRequest', '__init__', 'Error finding available receive mode directory') + self.upload_error = True + break + except PermissionError: + self.web.add_request(self.web.REQUEST_ERROR_DATA_DIR_CANNOT_CREATE, request.path, { + "receive_mode_dir": self.receive_mode_dir + }) + print(strings._('error_cannot_create_data_dir').format(self.receive_mode_dir)) + self.web.common.log('ReceiveModeRequest', '__init__', 'Permission denied creating receive mode directory') + self.upload_error = True + + # If there's an error so far, finish early + if self.upload_error: + return + # A dictionary that maps filenames to the bytes uploaded so far self.progress = {} @@ -331,7 +321,11 @@ class ReceiveModeRequest(Request): 'complete': False } - return ReceiveModeTemporaryFile(self, filename, self.file_write_func, self.file_close_func) + f = ReceiveModeFile(self, filename, self.file_write_func, self.file_close_func) + if f.upload_error: + self.web.common.log('ReceiveModeRequest', '_get_file_stream', 'Error creating file') + self.upload_error = True + return f def close(self): """ @@ -376,8 +370,6 @@ class ReceiveModeRequest(Request): self.progress[filename]['uploaded_bytes'] += length if self.previous_file != filename: - if self.previous_file is not None: - print('') self.previous_file = filename print('\r=> {:15s} {}'.format( diff --git a/onionshare/web/web.py b/onionshare/web/web.py index e2b22c4d..183c512c 100644 --- a/onionshare/web/web.py +++ b/onionshare/web/web.py @@ -14,7 +14,7 @@ from flask import Flask, request, render_template, abort, make_response, __versi from .. import strings from .share_mode import ShareModeWeb -from .receive_mode import ReceiveModeWeb, ReceiveModeWSGIMiddleware, ReceiveModeTemporaryFile, ReceiveModeRequest +from .receive_mode import ReceiveModeWeb, ReceiveModeWSGIMiddleware, ReceiveModeRequest # Stub out flask's show_server_banner function, to avoiding showing warnings that -- cgit v1.2.3-54-g00ecf