Refactor web even more to all of the share and receive web logic into ShareModeWeb and ReceiveModeWeb classes

This commit is contained in:
Micah Lee 2018-09-21 11:14:32 -07:00
parent 8ce90fdd60
commit cc9f646f8b
No known key found for this signature in database
GPG Key ID: 403C2657CD994F73
8 changed files with 387 additions and 366 deletions

View File

@ -119,7 +119,7 @@ def main(cwd=None):
# Prepare files to share # Prepare files to share
print(strings._("preparing_files")) print(strings._("preparing_files"))
try: try:
web.set_file_info(filenames) web.share_mode.set_file_info(filenames)
if web.is_zipped: if web.is_zipped:
app.cleanup_filenames.append(web.download_filename) app.cleanup_filenames.append(web.download_filename)
except OSError as e: except OSError as e:

View File

@ -8,156 +8,164 @@ from ..common import DownloadsDirErrorCannotCreate, DownloadsDirErrorNotWritable
from .. import strings from .. import strings
def receive_routes(web): class ReceiveModeWeb(object):
""" """
The web app routes for receiving files All of the web logic for receive mode
""" """
def index_logic(): def __init__(self, web):
web.add_request(web.REQUEST_LOAD, request.path) self.web = web
self.define_routes()
if web.common.settings.get('public_mode'): def define_routes(self):
upload_action = '/upload'
close_action = '/close'
else:
upload_action = '/{}/upload'.format(web.slug)
close_action = '/{}/close'.format(web.slug)
r = make_response(render_template(
'receive.html',
upload_action=upload_action,
close_action=close_action,
receive_allow_receiver_shutdown=web.common.settings.get('receive_allow_receiver_shutdown')))
return web.add_security_headers(r)
@web.app.route("/<slug_candidate>")
def index(slug_candidate):
web.check_slug_candidate(slug_candidate)
return index_logic()
@web.app.route("/")
def index_public():
if not web.common.settings.get('public_mode'):
return web.error404()
return index_logic()
def upload_logic(slug_candidate=''):
""" """
Upload files. The web app routes for receiving files
""" """
# Make sure downloads_dir exists def index_logic():
valid = True self.web.add_request(self.web.REQUEST_LOAD, request.path)
try:
web.common.validate_downloads_dir() if self.web.common.settings.get('public_mode'):
except DownloadsDirErrorCannotCreate: upload_action = '/upload'
web.add_request(web.REQUEST_ERROR_DOWNLOADS_DIR_CANNOT_CREATE, request.path) close_action = '/close'
print(strings._('error_cannot_create_downloads_dir').format(web.common.settings.get('downloads_dir'))) else:
valid = False upload_action = '/{}/upload'.format(self.web.slug)
except DownloadsDirErrorNotWritable: close_action = '/{}/close'.format(self.web.slug)
web.add_request(web.REQUEST_ERROR_DOWNLOADS_DIR_NOT_WRITABLE, request.path)
print(strings._('error_downloads_dir_not_writable').format(web.common.settings.get('downloads_dir'))) r = make_response(render_template(
valid = False 'receive.html',
if not valid: upload_action=upload_action,
flash('Error uploading, please inform the OnionShare user', 'error') close_action=close_action,
if web.common.settings.get('public_mode'): receive_allow_receiver_shutdown=self.web.common.settings.get('receive_allow_receiver_shutdown')))
return self.web.add_security_headers(r)
@self.web.app.route("/<slug_candidate>")
def index(slug_candidate):
self.web.check_slug_candidate(slug_candidate)
return index_logic()
@self.web.app.route("/")
def index_public():
if not self.web.common.settings.get('public_mode'):
return self.web.error404()
return index_logic()
def upload_logic(slug_candidate=''):
"""
Upload files.
"""
# Make sure downloads_dir exists
valid = True
try:
self.web.common.validate_downloads_dir()
except DownloadsDirErrorCannotCreate:
self.web.add_request(self.web.REQUEST_ERROR_DOWNLOADS_DIR_CANNOT_CREATE, request.path)
print(strings._('error_cannot_create_downloads_dir').format(self.web.common.settings.get('downloads_dir')))
valid = False
except DownloadsDirErrorNotWritable:
self.web.add_request(self.web.REQUEST_ERROR_DOWNLOADS_DIR_NOT_WRITABLE, request.path)
print(strings._('error_downloads_dir_not_writable').format(self.web.common.settings.get('downloads_dir')))
valid = False
if not valid:
flash('Error uploading, please inform the OnionShare user', 'error')
if self.web.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(self.web.common.settings.get('downloads_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(self.web.common.settings.get('downloads_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(self.web.common.settings.get('downloads_dir'), new_filename)
if os.path.exists(local_path):
i += 1
else:
valid = True
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
})
self.web.common.log('Web', 'receive_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,
# to avoid leaking the locale of the OnionShare user
if len(filenames) == 0:
flash('No files uploaded', 'info')
else:
for filename in filenames:
flash('Sent {}'.format(filename), 'info')
if self.web.common.settings.get('public_mode'):
return redirect('/') return redirect('/')
else: else:
return redirect('/{}'.format(slug_candidate)) return redirect('/{}'.format(slug_candidate))
files = request.files.getlist('file[]') @self.web.app.route("/<slug_candidate>/upload", methods=['POST'])
filenames = [] def upload(slug_candidate):
print('') self.web.check_slug_candidate(slug_candidate)
for f in files: return upload_logic(slug_candidate)
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(web.common.settings.get('downloads_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 @self.web.app.route("/upload", methods=['POST'])
valid = False def upload_public():
while not valid: if not self.web.common.settings.get('public_mode'):
new_filename = '{}-{}.{}'.format('.'.join(name), i, ext) return self.web.error404()
local_path = os.path.join(web.common.settings.get('downloads_dir'), new_filename) return upload_logic()
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(web.common.settings.get('downloads_dir'), new_filename)
if os.path.exists(local_path):
i += 1
else:
valid = True
basename = os.path.basename(local_path)
if f.filename != basename:
# Tell the GUI that the file has changed names
web.add_request(web.REQUEST_UPLOAD_FILE_RENAMED, request.path, {
'id': request.upload_id,
'old_filename': f.filename,
'new_filename': basename
})
web.common.log('Web', 'receive_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,
# to avoid leaking the locale of the OnionShare user
if len(filenames) == 0:
flash('No files uploaded', 'info')
else:
for filename in filenames:
flash('Sent {}'.format(filename), 'info')
if web.common.settings.get('public_mode'):
return redirect('/')
else:
return redirect('/{}'.format(slug_candidate))
@web.app.route("/<slug_candidate>/upload", methods=['POST'])
def upload(slug_candidate):
web.check_slug_candidate(slug_candidate)
return upload_logic(slug_candidate)
@web.app.route("/upload", methods=['POST'])
def upload_public():
if not web.common.settings.get('public_mode'):
return web.error404()
return upload_logic()
def close_logic(slug_candidate=''): def close_logic(slug_candidate=''):
if web.common.settings.get('receive_allow_receiver_shutdown'): if self.web.common.settings.get('receive_allow_receiver_shutdown'):
web.force_shutdown() self.web.force_shutdown()
r = make_response(render_template('closed.html')) r = make_response(render_template('closed.html'))
web.add_request(web.REQUEST_CLOSE_SERVER, request.path) self.web.add_request(self.web.REQUEST_CLOSE_SERVER, request.path)
return web.add_security_headers(r) return self.web.add_security_headers(r)
else: else:
return redirect('/{}'.format(slug_candidate)) return redirect('/{}'.format(slug_candidate))
@web.app.route("/<slug_candidate>/close", methods=['POST']) @self.web.app.route("/<slug_candidate>/close", methods=['POST'])
def close(slug_candidate): def close(slug_candidate):
web.check_slug_candidate(slug_candidate) self.web.check_slug_candidate(slug_candidate)
return close_logic(slug_candidate) return close_logic(slug_candidate)
@web.app.route("/close", methods=['POST']) @self.web.app.route("/close", methods=['POST'])
def close_public(): def close_public():
if not web.common.settings.get('public_mode'): if not self.web.common.settings.get('public_mode'):
return web.error404() return self.web.error404()
return close_logic() return close_logic()
class ReceiveModeWSGIMiddleware(object): class ReceiveModeWSGIMiddleware(object):

View File

@ -8,176 +8,237 @@ from flask import Response, request, render_template, make_response
from .. import strings from .. import strings
def share_routes(web): class ShareModeWeb(object):
""" """
The web app routes for sharing files All of the web logic for share mode
""" """
@web.app.route("/<slug_candidate>") def __init__(self, web):
def index(slug_candidate): self.web = web
web.check_slug_candidate(slug_candidate) self.define_routes()
return index_logic()
@web.app.route("/") def define_routes(self):
def index_public():
if not web.common.settings.get('public_mode'):
return web.error404()
return index_logic()
def index_logic(slug_candidate=''):
""" """
Render the template for the onionshare landing page. The web app routes for sharing files
""" """
web.add_request(web.REQUEST_LOAD, request.path) @self.web.app.route("/<slug_candidate>")
def index(slug_candidate):
self.web.check_slug_candidate(slug_candidate)
return index_logic()
# Deny new downloads if "Stop After First Download" is checked and there is @self.web.app.route("/")
# currently a download def index_public():
deny_download = not web.stay_open and web.download_in_progress if not self.web.common.settings.get('public_mode'):
if deny_download: return self.web.error404()
r = make_response(render_template('denied.html')) return index_logic()
return web.add_security_headers(r)
def index_logic(slug_candidate=''):
"""
Render the template for the onionshare landing page.
"""
self.web.add_request(self.web.REQUEST_LOAD, request.path)
# 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
if deny_download:
r = make_response(render_template('denied.html'))
return self.web.add_security_headers(r)
# If download is allowed to continue, serve download page
if web.slug:
r = make_response(render_template(
'send.html',
slug=web.slug,
file_info=web.file_info,
filename=os.path.basename(web.download_filename),
filesize=web.download_filesize,
filesize_human=web.common.human_readable_filesize(web.download_filesize),
is_zipped=web.is_zipped))
else:
# If download is allowed to continue, serve download page # If download is allowed to continue, serve download page
r = make_response(render_template( if self.web.slug:
'send.html', r = make_response(render_template(
file_info=web.file_info, 'send.html',
filename=os.path.basename(web.download_filename), slug=self.web.slug,
filesize=web.download_filesize, file_info=self.web.file_info,
filesize_human=web.common.human_readable_filesize(web.download_filesize), filename=os.path.basename(self.web.download_filename),
is_zipped=web.is_zipped)) filesize=self.web.download_filesize,
return web.add_security_headers(r) filesize_human=self.web.common.human_readable_filesize(self.web.download_filesize),
is_zipped=self.web.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))
return self.web.add_security_headers(r)
@web.app.route("/<slug_candidate>/download") @self.web.app.route("/<slug_candidate>/download")
def download(slug_candidate): def download(slug_candidate):
web.check_slug_candidate(slug_candidate) self.web.check_slug_candidate(slug_candidate)
return download_logic() return download_logic()
@web.app.route("/download") @self.web.app.route("/download")
def download_public(): def download_public():
if not web.common.settings.get('public_mode'): if not self.web.common.settings.get('public_mode'):
return web.error404() return self.web.error404()
return download_logic() return download_logic()
def download_logic(slug_candidate=''): def download_logic(slug_candidate=''):
""" """
Download the zip file. Download the zip file.
""" """
# Deny new downloads if "Stop After First Download" is checked and there is # Deny new downloads if "Stop After First Download" is checked and there is
# currently a download # currently a download
deny_download = not web.stay_open and web.download_in_progress deny_download = not self.web.stay_open and self.web.download_in_progress
if deny_download: if deny_download:
r = make_response(render_template('denied.html')) r = make_response(render_template('denied.html'))
return web.add_security_headers(r) return self.web.add_security_headers(r)
# Each download has a unique id # Each download has a unique id
download_id = web.download_count download_id = self.web.download_count
web.download_count += 1 self.web.download_count += 1
# Prepare some variables to use inside generate() function below # Prepare some variables to use inside generate() function below
# which is outside of the request context # which is outside of the request context
shutdown_func = request.environ.get('werkzeug.server.shutdown') shutdown_func = request.environ.get('werkzeug.server.shutdown')
path = request.path path = request.path
# Tell GUI the download started # Tell GUI the download started
web.add_request(web.REQUEST_STARTED, path, { self.web.add_request(self.web.REQUEST_STARTED, path, {
'id': download_id} 'id': download_id}
) )
dirname = os.path.dirname(web.download_filename) dirname = os.path.dirname(self.web.download_filename)
basename = os.path.basename(web.download_filename) basename = os.path.basename(self.web.download_filename)
def generate(): def generate():
# The user hasn't canceled the download # The user hasn't canceled the download
web.client_cancel = False self.web.client_cancel = False
# Starting a new download # Starting a new download
if not web.stay_open: if not self.web.stay_open:
web.download_in_progress = True self.web.download_in_progress = True
chunk_size = 102400 # 100kb chunk_size = 102400 # 100kb
fp = open(web.download_filename, 'rb') fp = open(self.web.download_filename, 'rb')
web.done = False self.web.done = False
canceled = False canceled = False
while not web.done: while not self.web.done:
# The user has canceled the download, so stop serving the file # The user has canceled the download, so stop serving the file
if web.client_cancel: if self.web.client_cancel:
web.add_request(web.REQUEST_CANCELED, path, { self.web.add_request(self.web.REQUEST_CANCELED, path, {
'id': download_id
})
break
chunk = fp.read(chunk_size)
if chunk == b'':
web.done = True
else:
try:
yield chunk
# tell GUI the progress
downloaded_bytes = fp.tell()
percent = (1.0 * downloaded_bytes / web.download_filesize) * 100
# only output to stdout if running onionshare in CLI mode, or if using Linux (#203, #304)
if not web.gui_mode or web.common.platform == 'Linux' or web.common.platform == 'BSD':
sys.stdout.write(
"\r{0:s}, {1:.2f}% ".format(web.common.human_readable_filesize(downloaded_bytes), percent))
sys.stdout.flush()
web.add_request(web.REQUEST_PROGRESS, path, {
'id': download_id,
'bytes': downloaded_bytes
})
web.done = False
except:
# looks like the download was canceled
web.done = True
canceled = True
# tell the GUI the download has canceled
web.add_request(web.REQUEST_CANCELED, path, {
'id': download_id 'id': download_id
}) })
break
fp.close() chunk = fp.read(chunk_size)
if chunk == b'':
self.web.done = True
else:
try:
yield chunk
if web.common.platform != 'Darwin': # tell GUI the progress
sys.stdout.write("\n") downloaded_bytes = fp.tell()
percent = (1.0 * downloaded_bytes / self.web.download_filesize) * 100
# Download is finished # only output to stdout if running onionshare in CLI mode, or if using Linux (#203, #304)
if not web.stay_open: if not self.web.is_gui or self.web.common.platform == 'Linux' or self.web.common.platform == 'BSD':
web.download_in_progress = False sys.stdout.write(
"\r{0:s}, {1:.2f}% ".format(self.web.common.human_readable_filesize(downloaded_bytes), percent))
sys.stdout.flush()
# Close the server, if necessary self.web.add_request(self.web.REQUEST_PROGRESS, path, {
if not web.stay_open and not canceled: 'id': download_id,
print(strings._("closing_automatically")) 'bytes': downloaded_bytes
web.running = False })
try: self.web.done = False
if shutdown_func is None: except:
raise RuntimeError('Not running with the Werkzeug Server') # looks like the download was canceled
shutdown_func() self.web.done = True
except: canceled = True
pass
r = Response(generate()) # tell the GUI the download has canceled
r.headers.set('Content-Length', web.download_filesize) self.web.add_request(self.web.REQUEST_CANCELED, path, {
r.headers.set('Content-Disposition', 'attachment', filename=basename) 'id': download_id
r = web.add_security_headers(r) })
# guess content type
(content_type, _) = mimetypes.guess_type(basename, strict=False) fp.close()
if content_type is not None:
r.headers.set('Content-Type', content_type) if self.web.common.platform != 'Darwin':
return r sys.stdout.write("\n")
# Download is finished
if not self.web.stay_open:
self.web.download_in_progress = False
# Close the server, if necessary
if not self.web.stay_open and not canceled:
print(strings._("closing_automatically"))
self.web.running = False
try:
if shutdown_func is None:
raise RuntimeError('Not running with the Werkzeug Server')
shutdown_func()
except:
pass
r = Response(generate())
r.headers.set('Content-Length', self.web.download_filesize)
r.headers.set('Content-Disposition', 'attachment', filename=basename)
r = self.web.add_security_headers(r)
# guess content type
(content_type, _) = mimetypes.guess_type(basename, strict=False)
if content_type is not None:
r.headers.set('Content-Type', content_type)
return r
def set_file_info(self, filenames, processed_size_callback=None):
"""
Using the list of filenames being shared, fill in details that the web
page will need to display. This includes zipping up the file in order to
get the zip file's name and size.
"""
self.web.common.log("Web", "set_file_info")
self.web.cancel_compression = False
# build file info list
self.web.file_info = {'files': [], 'dirs': []}
for filename in filenames:
info = {
'filename': filename,
'basename': os.path.basename(filename.rstrip('/'))
}
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)
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'])
# 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']
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'])
# Canceling early?
if self.web.cancel_compression:
self.web.zip_writer.close()
return False
for info in self.web.file_info['dirs']:
if not self.web.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
return True
class ZipWriter(object): class ZipWriter(object):

View File

@ -13,8 +13,8 @@ from flask import Flask, request, render_template, abort, make_response, __versi
from .. import strings from .. import strings
from .share_mode import share_routes, ZipWriter from .share_mode import ShareModeWeb
from .receive_mode import receive_routes, ReceiveModeWSGIMiddleware, ReceiveModeTemporaryFile, ReceiveModeRequest from .receive_mode import ReceiveModeWeb, ReceiveModeWSGIMiddleware, ReceiveModeTemporaryFile, ReceiveModeRequest
# Stub out flask's show_server_banner function, to avoiding showing warnings that # Stub out flask's show_server_banner function, to avoiding showing warnings that
@ -41,7 +41,7 @@ class Web(object):
REQUEST_ERROR_DOWNLOADS_DIR_CANNOT_CREATE = 9 REQUEST_ERROR_DOWNLOADS_DIR_CANNOT_CREATE = 9
REQUEST_ERROR_DOWNLOADS_DIR_NOT_WRITABLE = 10 REQUEST_ERROR_DOWNLOADS_DIR_NOT_WRITABLE = 10
def __init__(self, common, gui_mode, receive_mode=False): def __init__(self, common, is_gui, mode='share'):
self.common = common self.common = common
# The flask app # The flask app
@ -55,11 +55,11 @@ class Web(object):
self.debug_mode() self.debug_mode()
# Are we running in GUI mode? # Are we running in GUI mode?
self.gui_mode = gui_mode self.is_gui = is_gui
# Are we using receive mode? # Are we using receive mode?
self.receive_mode = receive_mode self.mode = mode
if self.receive_mode: if self.mode == 'receive':
# Use custom WSGI middleware, to modify environ # Use custom WSGI middleware, to modify environ
self.app.wsgi_app = ReceiveModeWSGIMiddleware(self.app.wsgi_app, self) self.app.wsgi_app = ReceiveModeWSGIMiddleware(self.app.wsgi_app, self)
# Use a custom Request class to track upload progess # Use a custom Request class to track upload progess
@ -115,14 +115,19 @@ class Web(object):
# Keep track if the server is running # Keep track if the server is running
self.running = False self.running = False
# Define the ewb app routes # Define the web app routes
self.common_routes() self.define_common_routes()
if self.receive_mode:
receive_routes(self)
else:
share_routes(self)
def common_routes(self): # Create the mode web object, which defines its own routes
self.share_mode = None
self.receive_mode = None
if self.mode == 'receive':
self.receive_mode = ReceiveModeWeb(self)
elif self.mode == 'share':
self.share_mode = ShareModeWeb(self)
def define_common_routes(self):
""" """
Common web app routes between sending and receiving Common web app routes between sending and receiving
""" """
@ -165,59 +170,6 @@ class Web(object):
r.headers.set(header, value) r.headers.set(header, value)
return r return r
def set_file_info(self, filenames, processed_size_callback=None):
"""
Using the list of filenames being shared, fill in details that the web
page will need to display. This includes zipping up the file in order to
get the zip file's name and size.
"""
self.common.log("Web", "set_file_info")
self.cancel_compression = False
# build file info list
self.file_info = {'files': [], 'dirs': []}
for filename in filenames:
info = {
'filename': filename,
'basename': os.path.basename(filename.rstrip('/'))
}
if os.path.isfile(filename):
info['size'] = os.path.getsize(filename)
info['size_human'] = self.common.human_readable_filesize(info['size'])
self.file_info['files'].append(info)
if os.path.isdir(filename):
info['size'] = self.common.dir_size(filename)
info['size_human'] = self.common.human_readable_filesize(info['size'])
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.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.zip_writer = ZipWriter(self.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.cancel_compression:
self.zip_writer.close()
return False
for info in self.file_info['dirs']:
if not self.zip_writer.add_dir(info['filename']):
return False
self.zip_writer.close()
self.download_filesize = os.path.getsize(self.download_filename)
self.is_zipped = True
return True
def _safe_select_jinja_autoescape(self, filename): def _safe_select_jinja_autoescape(self, filename):
if filename is None: if filename is None:
return True return True

View File

@ -34,7 +34,7 @@ class ReceiveMode(Mode):
Custom initialization for ReceiveMode. Custom initialization for ReceiveMode.
""" """
# Create the Web object # Create the Web object
self.web = Web(self.common, True, True) self.web = Web(self.common, True, 'receive')
# Server status # Server status
self.server_status.set_mode('receive') self.server_status.set_mode('receive')

View File

@ -43,7 +43,7 @@ class ShareMode(Mode):
self.compress_thread = None self.compress_thread = None
# Create the Web object # Create the Web object
self.web = Web(self.common, True, False) self.web = Web(self.common, True, 'share')
# File selection # File selection
self.file_selection = FileSelection(self.common) self.file_selection = FileSelection(self.common)

View File

@ -41,7 +41,7 @@ class CompressThread(QtCore.QThread):
self.mode.common.log('CompressThread', 'run') self.mode.common.log('CompressThread', 'run')
try: try:
if self.mode.web.set_file_info(self.mode.filenames, processed_size_callback=self.set_processed_size): if self.mode.web.share_mode.set_file_info(self.mode.filenames, processed_size_callback=self.set_processed_size):
self.success.emit() self.success.emit()
else: else:
# Cancelled # Cancelled

View File

@ -38,11 +38,11 @@ DEFAULT_ZW_FILENAME_REGEX = re.compile(r'^onionshare_[a-z2-7]{6}.zip$')
RANDOM_STR_REGEX = re.compile(r'^[a-z2-7]+$') RANDOM_STR_REGEX = re.compile(r'^[a-z2-7]+$')
def web_obj(common_obj, receive_mode, num_files=0): def web_obj(common_obj, mode, num_files=0):
""" Creates a Web object, in either share mode or receive mode, ready for testing """ """ Creates a Web object, in either share mode or receive mode, ready for testing """
common_obj.load_settings() common_obj.load_settings()
web = Web(common_obj, False, receive_mode) web = Web(common_obj, False, mode)
web.generate_slug() web.generate_slug()
web.stay_open = True web.stay_open = True
web.running = True web.running = True
@ -50,14 +50,14 @@ def web_obj(common_obj, receive_mode, num_files=0):
web.app.testing = True web.app.testing = True
# Share mode # Share mode
if not receive_mode: if mode == 'share':
# Add files # Add files
files = [] files = []
for i in range(num_files): for i in range(num_files):
with tempfile.NamedTemporaryFile(delete=False) as tmp_file: with tempfile.NamedTemporaryFile(delete=False) as tmp_file:
tmp_file.write(b'*' * 1024) tmp_file.write(b'*' * 1024)
files.append(tmp_file.name) files.append(tmp_file.name)
web.set_file_info(files) web.share_mode.set_file_info(files)
# Receive mode # Receive mode
else: else:
pass pass
@ -67,8 +67,8 @@ def web_obj(common_obj, receive_mode, num_files=0):
class TestWeb: class TestWeb:
def test_share_mode(self, common_obj): def test_share_mode(self, common_obj):
web = web_obj(common_obj, False, 3) web = web_obj(common_obj, 'share', 3)
assert web.receive_mode is False assert web.mode is 'share'
with web.app.test_client() as c: with web.app.test_client() as c:
# Load 404 pages # Load 404 pages
res = c.get('/') res = c.get('/')
@ -91,7 +91,7 @@ class TestWeb:
assert res.mimetype == 'application/zip' assert res.mimetype == 'application/zip'
def test_share_mode_close_after_first_download_on(self, common_obj, temp_file_1024): def test_share_mode_close_after_first_download_on(self, common_obj, temp_file_1024):
web = web_obj(common_obj, False, 3) web = web_obj(common_obj, 'share', 3)
web.stay_open = False web.stay_open = False
assert web.running == True assert web.running == True
@ -106,7 +106,7 @@ class TestWeb:
assert web.running == False assert web.running == False
def test_share_mode_close_after_first_download_off(self, common_obj, temp_file_1024): def test_share_mode_close_after_first_download_off(self, common_obj, temp_file_1024):
web = web_obj(common_obj, False, 3) web = web_obj(common_obj, 'share', 3)
web.stay_open = True web.stay_open = True
assert web.running == True assert web.running == True
@ -120,8 +120,8 @@ class TestWeb:
assert web.running == True assert web.running == True
def test_receive_mode(self, common_obj): def test_receive_mode(self, common_obj):
web = web_obj(common_obj, True) web = web_obj(common_obj, 'receive')
assert web.receive_mode is True assert web.mode is 'receive'
with web.app.test_client() as c: with web.app.test_client() as c:
# Load 404 pages # Load 404 pages
@ -139,7 +139,7 @@ class TestWeb:
assert res.status_code == 200 assert res.status_code == 200
def test_receive_mode_allow_receiver_shutdown_on(self, common_obj): def test_receive_mode_allow_receiver_shutdown_on(self, common_obj):
web = web_obj(common_obj, True) web = web_obj(common_obj, 'receive')
common_obj.settings.set('receive_allow_receiver_shutdown', True) common_obj.settings.set('receive_allow_receiver_shutdown', True)
@ -154,7 +154,7 @@ class TestWeb:
assert web.running == False assert web.running == False
def test_receive_mode_allow_receiver_shutdown_off(self, common_obj): def test_receive_mode_allow_receiver_shutdown_off(self, common_obj):
web = web_obj(common_obj, True) web = web_obj(common_obj, 'receive')
common_obj.settings.set('receive_allow_receiver_shutdown', False) common_obj.settings.set('receive_allow_receiver_shutdown', False)
@ -167,9 +167,9 @@ class TestWeb:
# Should redirect to index, and server should still be running # Should redirect to index, and server should still be running
assert res.status_code == 302 assert res.status_code == 302
assert web.running == True assert web.running == True
def test_public_mode_on(self, common_obj): def test_public_mode_on(self, common_obj):
web = web_obj(common_obj, True) web = web_obj(common_obj, 'receive')
common_obj.settings.set('public_mode', True) common_obj.settings.set('public_mode', True)
with web.app.test_client() as c: with web.app.test_client() as c:
@ -182,9 +182,9 @@ class TestWeb:
res = c.get('/{}'.format(web.slug)) res = c.get('/{}'.format(web.slug))
data2 = res.get_data() data2 = res.get_data()
assert res.status_code == 404 assert res.status_code == 404
def test_public_mode_off(self, common_obj): def test_public_mode_off(self, common_obj):
web = web_obj(common_obj, True) web = web_obj(common_obj, 'receive')
common_obj.settings.set('public_mode', False) common_obj.settings.set('public_mode', False)
with web.app.test_client() as c: with web.app.test_client() as c: