Refactor web.py to move all the web logic into the Web class, and refactor onionshare (cli) to work with it -- but onionshare_gui is currently broken

onionshare/__init__.py

@@ -20,7 +20,8 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 import os, sys, time, argparse, threading
 
-from . import strings, common, web
+from . import strings, common
+from .web import Web
 from .onion import *
 from .onionshare import OnionShare
 from .settings import Settings
@@ -67,14 +68,9 @@ def main(cwd=None):
         print(strings._('no_filenames'))
         sys.exit()
 
-    # Tell web if receive mode is enabled
-    if receive:
-        web.set_receive_mode()
-
     # Debug mode?
     if debug:
         common.set_debug(debug)
-        web.debug_mode()
 
     # Validation
     valid = True
@@ -88,10 +84,13 @@ def main(cwd=None):
     if not valid:
         sys.exit()
 
-
+    # Load settings
     settings = Settings(config)
     settings.load()
 
+    # Create the Web object
+    web = Web(debug, stay_open, False, receive)
+
     # Start the Onion object
     onion = Onion()
     try:

onionshare/web.py

@@ -37,11 +37,27 @@ from flask import (
 
 from . import strings, common
 
+class Web(object):
+    """
+    The Web object is the OnionShare web server, powered by flask
+    """
+    def __init__(self, debug, stay_open, gui_mode, receive_mode):
+        # The flask app
+        self.app = Flask(__name__)
+
+        # Debug mode?
+        if debug:
+            self.debug_mode()
+
+        # Stay open after the first download?
+        self.stay_open = False
+
+        # Are we running in GUI mode?
+        self.gui_mode = False
+
+        # Are we using receive mode?
+        self.receive_mode = False
 
-def _safe_select_jinja_autoescape(self, filename):
-    if filename is None:
-        return True
-    return filename.endswith(('.html', '.htm', '.xml', '.xhtml'))
 
         # Starting in Flask 0.11, render_template_string autoescapes template variables
         # by default. To prevent content injection through template variables in
@@ -49,16 +65,14 @@ def _safe_select_jinja_autoescape(self, filename):
         # engine if we detect a Flask version with insecure default behavior.
         if Version(flask_version) < Version('0.11'):
             # Monkey-patch in the fix from https://github.com/pallets/flask/commit/99c99c4c16b1327288fd76c44bc8635a1de452bc
-    Flask.select_jinja_autoescape = _safe_select_jinja_autoescape
+            Flask.select_jinja_autoescape = self._safe_select_jinja_autoescape
 
-app = Flask(__name__)
+        # Information about the file
+        self.file_info = []
+        self.zip_filename = None
+        self.zip_filesize = None
 
-# information about the file
+        self.security_headers = [
-file_info = []
-zip_filename = None
-zip_filesize = None
-
-security_headers = [
             ('Content-Security-Policy', 'default-src \'self\'; style-src \'unsafe-inline\'; script-src \'unsafe-inline\'; img-src \'self\' data:;'),
             ('X-Frame-Options', 'DENY'),
             ('X-Xss-Protection', '1; mode=block'),
@@ -67,222 +81,98 @@ security_headers = [
             ('Server', 'OnionShare')
         ]
 
-
+        self.REQUEST_LOAD = 0
-def set_file_info(filenames, processed_size_callback=None):
+        self.REQUEST_DOWNLOAD = 1
-    """
+        self.REQUEST_PROGRESS = 2
-    Using the list of filenames being shared, fill in details that the web
+        self.REQUEST_OTHER = 3
-    page will need to display. This includes zipping up the file in order to
+        self.REQUEST_CANCELED = 4
-    get the zip file's name and size.
+        self.REQUEST_RATE_LIMIT = 5
-    """
+        self.q = queue.Queue()
-    global file_info, zip_filename, zip_filesize
-
-    # build file info list
-    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'] = common.human_readable_filesize(info['size'])
-            file_info['files'].append(info)
-        if os.path.isdir(filename):
-            info['size'] = common.dir_size(filename)
-            info['size_human'] = common.human_readable_filesize(info['size'])
-            file_info['dirs'].append(info)
-    file_info['files'] = sorted(file_info['files'], key=lambda k: k['basename'])
-    file_info['dirs'] = sorted(file_info['dirs'], key=lambda k: k['basename'])
-
-    # zip up the files and folders
-    z = common.ZipWriter(processed_size_callback=processed_size_callback)
-    for info in file_info['files']:
-        z.add_file(info['filename'])
-    for info in file_info['dirs']:
-        z.add_dir(info['filename'])
-    z.close()
-    zip_filename = z.zip_filename
-    zip_filesize = os.path.getsize(zip_filename)
-
-
-REQUEST_LOAD = 0
-REQUEST_DOWNLOAD = 1
-REQUEST_PROGRESS = 2
-REQUEST_OTHER = 3
-REQUEST_CANCELED = 4
-REQUEST_RATE_LIMIT = 5
-q = queue.Queue()
-
-
-def add_request(request_type, path, data=None):
-    """
-    Add a request to the queue, to communicate with the GUI.
-    """
-    global q
-    q.put({
-        'type': request_type,
-        'path': path,
-        'data': data
-    })
-
 
         # Load and base64 encode images to pass into templates
-favicon_b64 = base64.b64encode(open(common.get_resource_path('images/favicon.ico'), 'rb').read()).decode()
+        self.favicon_b64 = self.base64_image('favicon.ico')
-logo_b64 = base64.b64encode(open(common.get_resource_path('images/logo.png'), 'rb').read()).decode()
+        self.logo_b64 = self.base64_image('logo.png')
-folder_b64 = base64.b64encode(open(common.get_resource_path('images/web_folder.png'), 'rb').read()).decode()
+        self.folder_b64 = self.base64_image('web_folder.png')
-file_b64 = base64.b64encode(open(common.get_resource_path('images/web_file.png'), 'rb').read()).decode()
+        self.file_b64 = self.base64_image('web_file.png')
 
-slug = None
+        self.slug = None
-
-
-def generate_slug(persistent_slug=''):
-    global slug
-    if persistent_slug:
-        slug = persistent_slug
-    else:
-        slug = common.build_slug()
-
-download_count = 0
-error404_count = 0
-
-stay_open = False
-
-
-def set_stay_open(new_stay_open):
-    """
-    Set stay_open variable.
-    """
-    global stay_open
-    stay_open = new_stay_open
-
-
-def get_stay_open():
-    """
-    Get stay_open variable.
-    """
-    return stay_open
-
-
-# Are we running in GUI mode?
-gui_mode = False
-
-
-def set_gui_mode():
-    """
-    Tell the web service that we're running in GUI mode
-    """
-    global gui_mode
-    gui_mode = True
-
-
-# Are we using receive mode?
-receive_mode = False
-
-
-def set_receive_mode():
-    """
-    Tell the web service that we're running in GUI mode
-    """
-    global receive_mode
-    receive_mode = True
-    print('receive mode enabled')
-
-
-def debug_mode():
-    """
-    Turn on debugging mode, which will log flask errors to a debug file.
-    """
-    temp_dir = tempfile.gettempdir()
-    log_handler = logging.FileHandler(
-        os.path.join(temp_dir, 'onionshare_server.log'))
-    log_handler.setLevel(logging.WARNING)
-    app.logger.addHandler(log_handler)
-
-
-def check_slug_candidate(slug_candidate, slug_compare=None):
-    if not slug_compare:
-        slug_compare = slug
-    if not hmac.compare_digest(slug_compare, slug_candidate):
-        abort(404)
 
+        self.download_count = 0
+        self.error404_count = 0
 
         # If "Stop After First Download" is checked (stay_open == False), only allow
         # one download at a time.
-download_in_progress = False
+        self.download_in_progress = False
-
-done = False
-
-@app.route("/<slug_candidate>")
-def index(slug_candidate):
-    """
-    Render the template for the onionshare landing page.
-    """
-    check_slug_candidate(slug_candidate)
-
-    add_request(REQUEST_LOAD, request.path)
-
-    # Deny new downloads if "Stop After First Download" is checked and there is
-    # currently a download
-    global stay_open, download_in_progress
-    deny_download = not stay_open and download_in_progress
-    if deny_download:
-        r = make_response(render_template_string(
-            open(common.get_resource_path('html/denied.html')).read(),
-            favicon_b64=favicon_b64
-        ))
-        for header, value in security_headers:
-            r.headers.set(header, value)
-        return r
-
-    # If download is allowed to continue, serve download page
-
-    r = make_response(render_template_string(
-        open(common.get_resource_path('html/index.html')).read(),
-        favicon_b64=favicon_b64,
-        logo_b64=logo_b64,
-        folder_b64=folder_b64,
-        file_b64=file_b64,
-        slug=slug,
-        file_info=file_info,
-        filename=os.path.basename(zip_filename),
-        filesize=zip_filesize,
-        filesize_human=common.human_readable_filesize(zip_filesize)))
-    for header, value in security_headers:
-        r.headers.set(header, value)
-    return r
 
+        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.
-client_cancel = False
+        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 = common.random_string(16)
 
-@app.route("/<slug_candidate>/download")
+        @self.app.route("/<slug_candidate>")
+        def index(slug_candidate):
+            """
+            Render the template for the onionshare landing page.
+            """
+            self.check_slug_candidate(slug_candidate)
+
+            self.add_request(self.REQUEST_LOAD, request.path)
+
+            # Deny new downloads if "Stop After First Download" is checked and there is
+            # currently a download
+            deny_download = not self.stay_open and self.download_in_progress
+            if deny_download:
+                r = make_response(render_template_string(
+                    open(common.get_resource_path('html/denied.html')).read(),
+                    favicon_b64=self.favicon_b64
+                ))
+                for header, value in self.security_headers:
+                    r.headers.set(header, value)
+                return r
+
+            # If download is allowed to continue, serve download page
+            r = make_response(render_template_string(
+                open(common.get_resource_path('html/index.html')).read(),
+                favicon_b64=self.favicon_b64,
+                logo_b64=self.logo_b64,
+                folder_b64=self.folder_b64,
+                file_b64=self.file_b64,
+                slug=self.slug,
+                file_info=self.file_info,
+                filename=os.path.basename(self.zip_filename),
+                filesize=self.zip_filesize,
+                filesize_human=common.human_readable_filesize(self.zip_filesize)))
+            for header, value in self.security_headers:
+                r.headers.set(header, value)
+            return r
+
+        @self.app.route("/<slug_candidate>/download")
         def download(slug_candidate):
             """
             Download the zip file.
             """
-    check_slug_candidate(slug_candidate)
+            self.check_slug_candidate(slug_candidate)
 
             # Deny new downloads if "Stop After First Download" is checked and there is
             # currently a download
-    global stay_open, download_in_progress, done
+            deny_download = not self.stay_open and self.download_in_progress
-    deny_download = not stay_open and download_in_progress
             if deny_download:
                 r = make_response(render_template_string(
                     open(common.get_resource_path('html/denied.html')).read(),
-            favicon_b64=favicon_b64
+                    favicon_b64=self.favicon_b64
                 ))
-        for header,value in security_headers:
+                for header,value in self.security_headers:
                     r.headers.set(header, value)
                 return r
 
-    global download_count
-
             # each download has a unique id
-    download_id = download_count
+            download_id = self.download_count
-    download_count += 1
+            self.download_count += 1
 
             # prepare some variables to use inside generate() function below
             # which is outside of the request context
@@ -290,59 +180,57 @@ def download(slug_candidate):
             path = request.path
 
             # tell GUI the download started
-    add_request(REQUEST_DOWNLOAD, path, {'id': download_id})
+            self.add_request(self.REQUEST_DOWNLOAD, path, {'id': download_id})
 
-    dirname = os.path.dirname(zip_filename)
+            dirname = os.path.dirname(self.zip_filename)
-    basename = os.path.basename(zip_filename)
+            basename = os.path.basename(self.zip_filename)
 
             def generate():
                 # The user hasn't canceled the download
-        global client_cancel, gui_mode
+                self.client_cancel = False
-        client_cancel = False
 
                 # Starting a new download
-        global stay_open, download_in_progress, done
+                if not self.stay_open:
-        if not stay_open:
+                    self.download_in_progress = True
-            download_in_progress = True
 
                 chunk_size = 102400  # 100kb
 
-        fp = open(zip_filename, 'rb')
+                fp = open(self.zip_filename, 'rb')
-        done = False
+                self.done = False
                 canceled = False
-        while not done:
+                while not self.done:
                     # The user has canceled the download, so stop serving the file
-            if client_cancel:
+                    if self.client_cancel:
-                add_request(REQUEST_CANCELED, path, {'id': download_id})
+                        self.add_request(self.REQUEST_CANCELED, path, {'id': download_id})
                         break
 
                     chunk = fp.read(chunk_size)
                     if chunk == b'':
-                done = True
+                        self.done = True
                     else:
                         try:
                             yield chunk
 
                             # tell GUI the progress
                             downloaded_bytes = fp.tell()
-                    percent = (1.0 * downloaded_bytes / zip_filesize) * 100
+                            percent = (1.0 * downloaded_bytes / self.zip_filesize) * 100
 
                             # only output to stdout if running onionshare in CLI mode, or if using Linux (#203, #304)
                             plat = common.get_platform()
-                    if not gui_mode or plat == 'Linux' or plat == 'BSD':
+                            if not self.gui_mode or plat == 'Linux' or plat == 'BSD':
                                 sys.stdout.write(
                                     "\r{0:s}, {1:.2f}%          ".format(common.human_readable_filesize(downloaded_bytes), percent))
                                 sys.stdout.flush()
 
-                    add_request(REQUEST_PROGRESS, path, {'id': download_id, 'bytes': downloaded_bytes})
+                            self.add_request(self.REQUEST_PROGRESS, path, {'id': download_id, 'bytes': downloaded_bytes})
-                    done = False
+                            self.done = False
                         except:
                             # looks like the download was canceled
-                    done = True
+                            self.done = True
                             canceled = True
 
                             # tell the GUI the download has canceled
-                    add_request(REQUEST_CANCELED, path, {'id': download_id})
+                            self.add_request(self.REQUEST_CANCELED, path, {'id': download_id})
 
                 fp.close()
 
@@ -350,20 +238,20 @@ def download(slug_candidate):
                     sys.stdout.write("\n")
 
                 # Download is finished
-        if not stay_open:
+                if not self.stay_open:
-            download_in_progress = False
+                    self.download_in_progress = False
 
                 # Close the server, if necessary
-        if not stay_open and not canceled:
+                if not self.stay_open and not canceled:
                     print(strings._("closing_automatically"))
                     if shutdown_func is None:
                         raise RuntimeError('Not running with the Werkzeug Server')
                     shutdown_func()
 
             r = Response(generate())
-    r.headers.set('Content-Length', zip_filesize)
+            r.headers.set('Content-Length', self.zip_filesize)
             r.headers.set('Content-Disposition', 'attachment', filename=basename)
-    for header,value in security_headers:
+            for header,value in self.security_headers:
                 r.headers.set(header, value)
             # guess content type
             (content_type, _) = mimetypes.guess_type(basename, strict=False)
@@ -371,36 +259,29 @@ def download(slug_candidate):
                 r.headers.set('Content-Type', content_type)
             return r
 
-
+        @self.app.errorhandler(404)
-@app.errorhandler(404)
         def page_not_found(e):
             """
             404 error page.
             """
-    add_request(REQUEST_OTHER, request.path)
+            self.add_request(self.REQUEST_OTHER, request.path)
 
-    global error404_count
             if request.path != '/favicon.ico':
-        error404_count += 1
+                self.error404_count += 1
-        if error404_count == 20:
+                if self.error404_count == 20:
-            add_request(REQUEST_RATE_LIMIT, request.path)
+                    self.add_request(self.REQUEST_RATE_LIMIT, request.path)
                     force_shutdown()
                     print(strings._('error_rate_limit'))
 
             r = make_response(render_template_string(
                 open(common.get_resource_path('html/404.html')).read(),
-        favicon_b64=favicon_b64
+                favicon_b64=self.favicon_b64
             ), 404)
-    for header, value in security_headers:
+            for header, value in self.security_headers:
                 r.headers.set(header, value)
             return r
 
-
+        @self.app.route("/<slug_candidate>/shutdown")
-# shutting down the server only works within the context of flask, so the easiest way to do it is over http
-shutdown_slug = common.random_string(16)
-
-
-@app.route("/<slug_candidate>/shutdown")
         def shutdown(slug_candidate):
             """
             Stop the flask web server, from the context of an http request.
@@ -409,8 +290,84 @@ def shutdown(slug_candidate):
             force_shutdown()
             return ""
 
+    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.
+        """
+        # 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'] = common.human_readable_filesize(info['size'])
+                self.file_info['files'].append(info)
+            if os.path.isdir(filename):
+                info['size'] = common.dir_size(filename)
+                info['size_human'] = 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'])
 
-def force_shutdown():
+        # zip up the files and folders
+        z = common.ZipWriter(processed_size_callback=processed_size_callback)
+        for info in self.file_info['files']:
+            z.add_file(info['filename'])
+        for info in self.file_info['dirs']:
+            z.add_dir(info['filename'])
+        z.close()
+        self.zip_filename = z.zip_filename
+        self.zip_filesize = os.path.getsize(self.zip_filename)
+
+    def _safe_select_jinja_autoescape(self, filename):
+        if filename is None:
+            return True
+        return filename.endswith(('.html', '.htm', '.xml', '.xhtml'))
+
+    def base64_image(self, filename):
+        """
+        Base64-encode an image file to use data URIs in the web app
+        """
+        return base64.b64encode(open(common.get_resource_path('images/{}'.format(filename)), 'rb').read()).decode()
+
+    def add_request(self, request_type, path, data=None):
+        """
+        Add a request to the queue, to communicate with the GUI.
+        """
+        self.q.put({
+            'type': request_type,
+            'path': path,
+            'data': data
+        })
+
+    def generate_slug(self, persistent_slug=''):
+        if persistent_slug:
+            self.slug = persistent_slug
+        else:
+            self.slug = common.build_slug()
+
+    def debug_mode(self):
+        """
+        Turn on debugging mode, which will log flask errors to a debug file.
+        """
+        temp_dir = tempfile.gettempdir()
+        log_handler = logging.FileHandler(
+            os.path.join(temp_dir, 'onionshare_server.log'))
+        log_handler.setLevel(logging.WARNING)
+        self.app.logger.addHandler(log_handler)
+
+    def check_slug_candidate(self, slug_candidate, slug_compare=None):
+        if not slug_compare:
+            slug_compare = self.slug
+        if not hmac.compare_digest(slug_compare, slug_candidate):
+            abort(404)
+
+    def force_shutdown(self):
         """
         Stop the flask web server, from the context of the flask app.
         """
@@ -420,14 +377,13 @@ def force_shutdown():
             raise RuntimeError('Not running with the Werkzeug Server')
         func()
 
-
+    def start(self, port, stay_open=False, persistent_slug=''):
-def start(port, stay_open=False, persistent_slug=''):
         """
         Start the flask web server.
         """
-    generate_slug(persistent_slug)
+        self.generate_slug(persistent_slug)
 
-    set_stay_open(stay_open)
+        self.stay_open = stay_open
 
         # In Whonix, listen on 0.0.0.0 instead of 127.0.0.1 (#220)
         if os.path.exists('/usr/share/anon-ws-base-files/workstation'):
@@ -435,18 +391,16 @@ def start(port, stay_open=False, persistent_slug=''):
         else:
             host = '127.0.0.1'
 
-    app.run(host=host, port=port, threaded=True)
+        self.app.run(host=host, port=port, threaded=True)
 
-
+    def stop(self, port):
-def stop(port):
         """
         Stop the flask web server by loading /shutdown.
         """
 
         # If the user cancels the download, let the download function know to stop
         # serving the file
-    global client_cancel
+        self.client_cancel = True
-    client_cancel = True
 
         # to stop flask, load http://127.0.0.1:<port>/<shutdown_slug>/shutdown
         try:

onionshare_gui/__init__.py

@@ -22,7 +22,8 @@ import os, sys, platform, argparse
 from .alert import Alert
 from PyQt5 import QtCore, QtWidgets
 
-from onionshare import strings, common, web
+from onionshare import strings, common
+from .web import Web
 from onionshare.onion import Onion
 from onionshare.onionshare import OnionShare
 from onionshare.settings import Settings