Move more mode-specific logic out of the Web class and into the approprate mode web classes

This commit is contained in:
Micah Lee 2018-09-21 11:36:19 -07:00
parent a86681e903
commit 28fd67cbcc
No known key found for this signature in database
GPG Key ID: 403C2657CD994F73
7 changed files with 87 additions and 84 deletions

View File

@ -126,14 +126,14 @@ def main(cwd=None):
print(strings._("preparing_files"))
try:
web.share_mode.set_file_info(filenames)
if web.is_zipped:
app.cleanup_filenames.append(web.download_filename)
if web.share_mode.is_zipped:
app.cleanup_filenames.append(web.share_mode.download_filename)
except OSError as e:
print(e.strerror)
sys.exit(1)
# Warn about sending large files over Tor
if web.download_filesize >= 157286400: # 150mb
if web.share_mode.download_filesize >= 157286400: # 150mb
print('')
print(strings._("large_filesize"))
print('')
@ -193,11 +193,12 @@ def main(cwd=None):
if app.shutdown_timeout > 0:
# if the shutdown timer was set and has run out, stop the server
if not app.shutdown_timer.is_alive():
# If there were no attempts to download the share, or all downloads are done, we can stop
if web.download_count == 0 or web.done:
print(strings._("close_on_timeout"))
web.stop(app.port)
break
if mode == 'share':
# If there were no attempts to download the share, or all downloads are done, we can stop
if web.share_mode.download_count == 0 or web.done:
print(strings._("close_on_timeout"))
web.stop(app.port)
break
# Allow KeyboardInterrupt exception to be handled with threads
# https://stackoverflow.com/questions/3788208/python-threading-ignores-keyboardinterrupt-exception
time.sleep(0.2)

View File

@ -14,6 +14,9 @@ class ReceiveModeWeb(object):
"""
def __init__(self, web):
self.web = web
self.upload_count = 0
self.define_routes()
def define_routes(self):
@ -243,8 +246,8 @@ class ReceiveModeRequest(Request):
self.progress = {}
# Create an upload_id, attach it to the request
self.upload_id = self.web.upload_count
self.web.upload_count += 1
self.upload_id = self.upload_count
self.upload_count += 1
# Figure out the content length
try:

View File

@ -14,6 +14,25 @@ class ShareModeWeb(object):
"""
def __init__(self, web):
self.web = web
# Information about the file to be shared
self.file_info = []
self.is_zipped = False
self.download_filename = None
self.download_filesize = None
self.zip_writer = None
self.download_count = 0
# If "Stop After First Download" is checked (stay_open == False), only allow
# one download at a time.
self.download_in_progress = False
# If the client closes the OnionShare window while a download is in progress,
# it should immediately stop serving the file. The client_cancel global is
# used to tell the download function that the client is canceling the download.
self.client_cancel = False
self.define_routes()
def define_routes(self):
@ -39,7 +58,7 @@ class ShareModeWeb(object):
# Deny new downloads if "Stop After First Download" is checked and there is
# currently a download
deny_download = not self.web.stay_open and self.web.download_in_progress
deny_download = not self.web.stay_open and self.download_in_progress
if deny_download:
r = make_response(render_template('denied.html'))
return self.web.add_security_headers(r)
@ -49,20 +68,20 @@ class ShareModeWeb(object):
r = make_response(render_template(
'send.html',
slug=self.web.slug,
file_info=self.web.file_info,
filename=os.path.basename(self.web.download_filename),
filesize=self.web.download_filesize,
filesize_human=self.web.common.human_readable_filesize(self.web.download_filesize),
is_zipped=self.web.is_zipped))
file_info=self.file_info,
filename=os.path.basename(self.download_filename),
filesize=self.download_filesize,
filesize_human=self.web.common.human_readable_filesize(self.download_filesize),
is_zipped=self.is_zipped))
else:
# If download is allowed to continue, serve download page
r = make_response(render_template(
'send.html',
file_info=self.web.file_info,
filename=os.path.basename(self.web.download_filename),
filesize=self.web.download_filesize,
filesize_human=self.web.common.human_readable_filesize(self.web.download_filesize),
is_zipped=self.web.is_zipped))
file_info=self.file_info,
filename=os.path.basename(self.download_filename),
filesize=self.download_filesize,
filesize_human=self.web.common.human_readable_filesize(self.download_filesize),
is_zipped=self.is_zipped))
return self.web.add_security_headers(r)
@self.web.app.route("/<slug_candidate>/download")
@ -82,14 +101,14 @@ class ShareModeWeb(object):
"""
# Deny new downloads if "Stop After First Download" is checked and there is
# currently a download
deny_download = not self.web.stay_open and self.web.download_in_progress
deny_download = not self.web.stay_open and self.download_in_progress
if deny_download:
r = make_response(render_template('denied.html'))
return self.web.add_security_headers(r)
# Each download has a unique id
download_id = self.web.download_count
self.web.download_count += 1
download_id = self.download_count
self.download_count += 1
# Prepare some variables to use inside generate() function below
# which is outside of the request context
@ -101,25 +120,25 @@ class ShareModeWeb(object):
'id': download_id}
)
dirname = os.path.dirname(self.web.download_filename)
basename = os.path.basename(self.web.download_filename)
dirname = os.path.dirname(self.download_filename)
basename = os.path.basename(self.download_filename)
def generate():
# The user hasn't canceled the download
self.web.client_cancel = False
self.client_cancel = False
# Starting a new download
if not self.web.stay_open:
self.web.download_in_progress = True
self.download_in_progress = True
chunk_size = 102400 # 100kb
fp = open(self.web.download_filename, 'rb')
fp = open(self.download_filename, 'rb')
self.web.done = False
canceled = False
while not self.web.done:
# The user has canceled the download, so stop serving the file
if self.web.client_cancel:
if self.client_cancel:
self.web.add_request(self.web.REQUEST_CANCELED, path, {
'id': download_id
})
@ -134,7 +153,7 @@ class ShareModeWeb(object):
# tell GUI the progress
downloaded_bytes = fp.tell()
percent = (1.0 * downloaded_bytes / self.web.download_filesize) * 100
percent = (1.0 * downloaded_bytes / self.download_filesize) * 100
# only output to stdout if running onionshare in CLI mode, or if using Linux (#203, #304)
if not self.web.is_gui or self.web.common.platform == 'Linux' or self.web.common.platform == 'BSD':
@ -164,7 +183,7 @@ class ShareModeWeb(object):
# Download is finished
if not self.web.stay_open:
self.web.download_in_progress = False
self.download_in_progress = False
# Close the server, if necessary
if not self.web.stay_open and not canceled:
@ -178,7 +197,7 @@ class ShareModeWeb(object):
pass
r = Response(generate())
r.headers.set('Content-Length', self.web.download_filesize)
r.headers.set('Content-Length', self.download_filesize)
r.headers.set('Content-Disposition', 'attachment', filename=basename)
r = self.web.add_security_headers(r)
# guess content type
@ -197,7 +216,7 @@ class ShareModeWeb(object):
self.web.cancel_compression = False
# build file info list
self.web.file_info = {'files': [], 'dirs': []}
self.file_info = {'files': [], 'dirs': []}
for filename in filenames:
info = {
'filename': filename,
@ -206,37 +225,37 @@ class ShareModeWeb(object):
if os.path.isfile(filename):
info['size'] = os.path.getsize(filename)
info['size_human'] = self.web.common.human_readable_filesize(info['size'])
self.web.file_info['files'].append(info)
self.file_info['files'].append(info)
if os.path.isdir(filename):
info['size'] = self.web.common.dir_size(filename)
info['size_human'] = self.web.common.human_readable_filesize(info['size'])
self.web.file_info['dirs'].append(info)
self.web.file_info['files'] = sorted(self.web.file_info['files'], key=lambda k: k['basename'])
self.web.file_info['dirs'] = sorted(self.web.file_info['dirs'], key=lambda k: k['basename'])
self.file_info['dirs'].append(info)
self.file_info['files'] = sorted(self.file_info['files'], key=lambda k: k['basename'])
self.file_info['dirs'] = sorted(self.file_info['dirs'], key=lambda k: k['basename'])
# Check if there's only 1 file and no folders
if len(self.web.file_info['files']) == 1 and len(self.web.file_info['dirs']) == 0:
self.web.is_zipped = False
self.web.download_filename = self.web.file_info['files'][0]['filename']
self.web.download_filesize = self.web.file_info['files'][0]['size']
if len(self.file_info['files']) == 1 and len(self.file_info['dirs']) == 0:
self.is_zipped = False
self.download_filename = self.file_info['files'][0]['filename']
self.download_filesize = self.file_info['files'][0]['size']
else:
# Zip up the files and folders
self.web.zip_writer = ZipWriter(self.web.common, processed_size_callback=processed_size_callback)
self.web.download_filename = self.web.zip_writer.zip_filename
for info in self.web.file_info['files']:
self.web.zip_writer.add_file(info['filename'])
self.zip_writer = ZipWriter(self.web.common, processed_size_callback=processed_size_callback)
self.download_filename = self.zip_writer.zip_filename
for info in self.file_info['files']:
self.zip_writer.add_file(info['filename'])
# Canceling early?
if self.web.cancel_compression:
self.web.zip_writer.close()
self.zip_writer.close()
return False
for info in self.web.file_info['dirs']:
if not self.web.zip_writer.add_dir(info['filename']):
for info in self.file_info['dirs']:
if not self.zip_writer.add_dir(info['filename']):
return False
self.web.zip_writer.close()
self.web.download_filesize = os.path.getsize(self.web.download_filename)
self.web.is_zipped = True
self.zip_writer.close()
self.download_filesize = os.path.getsize(self.download_filename)
self.is_zipped = True
return True

View File

@ -73,13 +73,6 @@ class Web(object):
# Monkey-patch in the fix from https://github.com/pallets/flask/commit/99c99c4c16b1327288fd76c44bc8635a1de452bc
Flask.select_jinja_autoescape = self._safe_select_jinja_autoescape
# Information about the file
self.file_info = []
self.is_zipped = False
self.download_filename = None
self.download_filesize = None
self.zip_writer = None
self.security_headers = [
('Content-Security-Policy', 'default-src \'self\'; style-src \'self\'; script-src \'self\'; img-src \'self\' data:;'),
('X-Frame-Options', 'DENY'),
@ -90,25 +83,11 @@ class Web(object):
]
self.q = queue.Queue()
self.slug = None
self.download_count = 0
self.upload_count = 0
self.error404_count = 0
# If "Stop After First Download" is checked (stay_open == False), only allow
# one download at a time.
self.download_in_progress = False
self.done = False
# If the client closes the OnionShare window while a download is in progress,
# it should immediately stop serving the file. The client_cancel global is
# used to tell the download function that the client is canceling the download.
self.client_cancel = False
# shutting down the server only works within the context of flask, so the easiest way to do it is over http
self.shutdown_slug = self.common.random_string(16)
@ -254,9 +233,10 @@ class Web(object):
Stop the flask web server by loading /shutdown.
"""
# If the user cancels the download, let the download function know to stop
# serving the file
self.client_cancel = True
if self.mode == 'share':
# If the user cancels the download, let the download function know to stop
# serving the file
self.share_mode.client_cancel = True
# To stop flask, load http://127.0.0.1:<port>/<shutdown_slug>/shutdown
if self.running:

View File

@ -100,7 +100,7 @@ class ReceiveMode(Mode):
Starting the server.
"""
# Reset web counters
self.web.upload_count = 0
self.web.receive_mode.upload_count = 0
self.web.error404_count = 0
# Hide and reset the uploads if we have previously shared

View File

@ -125,7 +125,7 @@ class ShareMode(Mode):
The shutdown timer expired, should we stop the server? Returns a bool
"""
# If there were no attempts to download the share, or all downloads are done, we can stop
if self.web.download_count == 0 or self.web.done:
if self.web.share_mode.download_count == 0 or self.web.done:
self.server_status.stop_server()
self.server_status_label.setText(strings._('close_on_timeout', True))
return True
@ -139,7 +139,7 @@ class ShareMode(Mode):
Starting the server.
"""
# Reset web counters
self.web.download_count = 0
self.web.share_mode.download_count = 0
self.web.error404_count = 0
# Hide and reset the downloads if we have previously shared
@ -177,7 +177,7 @@ class ShareMode(Mode):
self._zip_progress_bar = None
# Warn about sending large files over Tor
if self.web.download_filesize >= 157286400: # 150mb
if self.web.share_mode.download_filesize >= 157286400: # 150mb
self.filesize_warning.setText(strings._("large_filesize", True))
self.filesize_warning.show()
@ -229,7 +229,7 @@ class ShareMode(Mode):
"""
Handle REQUEST_STARTED event.
"""
self.downloads.add(event["data"]["id"], self.web.download_filesize)
self.downloads.add(event["data"]["id"], self.web.share_mode.download_filesize)
self.downloads_in_progress += 1
self.update_downloads_in_progress()
@ -242,7 +242,7 @@ class ShareMode(Mode):
self.downloads.update(event["data"]["id"], event["data"]["bytes"])
# Is the download complete?
if event["data"]["bytes"] == self.web.download_filesize:
if event["data"]["bytes"] == self.web.share_mode.download_filesize:
self.system_tray.showMessage(strings._('systray_download_completed_title', True), strings._('systray_download_completed_message', True))
# Update the total 'completed downloads' info

View File

@ -47,8 +47,8 @@ class CompressThread(QtCore.QThread):
# Cancelled
pass
if self.mode.web.is_zipped:
self.mode.app.cleanup_filenames.append(self.mode.web.download_filename)
if self.mode.web.share_mode.is_zipped:
self.mode.app.cleanup_filenames.append(self.mode.web.share_mode.download_filename)
except OSError as e:
self.error.emit(e.strerror)