Move all mode-specific code out of SendBaseModeWeb and into inherited methods in WebsiteModeWeb and ShareModeWeb

5 files changed, 151 insertions, 119 deletions

diff --git a/onionshare/web/send_base_mode.py b/onionshare/web/send_base_mode.py
index 0ca1b306..80f9e315 100644
--- a/onionshare/web/send_base_mode.py
+++ b/onionshare/web/send_base_mode.py
@@ -2,10 +2,11 @@ import os
 import sys
 import tempfile
 import mimetypes
-from flask import Response, request, render_template, make_response, send_from_directory
+from flask import Response, request, render_template, make_response
 
 from .. import strings
 
+
 class SendBaseModeWeb:
     """
     All of the web logic shared between share and website mode (modes where the user sends files)
@@ -40,44 +41,29 @@ class SendBaseModeWeb:
 
         self.define_routes()
 
-
     def init(self):
+        self.common.log('SendBaseModeWeb', '__init__')
+        self.define_routes()
+
+    def define_routes(self):
         """
-        Add custom initialization here.
+        Inherited class will implement this
         """
         pass
 
+    def directory_listing_template(self):
+        """
+        Inherited class will implement this. It should call render_template and return
+        the response.
+        """
+        pass
 
     def directory_listing(self, filenames, path='', filesystem_path=None):
         # If filesystem_path is None, this is the root directory listing
-        files = []
-        dirs = []
-        r = ''
-
         files, dirs = self.build_directory_listing(filenames, filesystem_path)
-
-        if self.web.mode == 'website':
-            r = make_response(render_template('listing.html',
-                path=path,
-                files=files,
-                dirs=dirs,
-                static_url_path=self.web.static_url_path))
-
-        elif self.web.mode == 'share':
-            r = make_response(render_template(
-                'send.html',
-                file_info=self.file_info,
-                files=files,
-                dirs=dirs,
-                filename=os.path.basename(self.download_filename),
-                filesize=self.filesize,
-                filesize_human=self.common.human_readable_filesize(self.download_filesize),
-                is_zipped=self.is_zipped,
-                static_url_path=self.web.static_url_path))
-
+        r = self.directory_listing_template(path, files, dirs)
         return self.web.add_security_headers(r)
 
-
     def build_directory_listing(self, filenames, filesystem_path):
         files = []
         dirs = []
@@ -103,6 +89,11 @@ class SendBaseModeWeb:
                 })
         return files, dirs
 
+    def set_file_info_custom(self, filenames, processed_size_callback):
+        """
+        Inherited class will implement this.
+        """
+        pass
 
     def set_file_info(self, filenames, processed_size_callback=None):
         """
@@ -114,18 +105,7 @@ class SendBaseModeWeb:
             filenames = [os.path.join(filenames[0], x) for x in os.listdir(filenames[0])]
 
         self.build_file_list(filenames)
-
-        if self.web.mode == 'share':
-            self.common.log("ShareModeWeb", "set_file_info")
-            self.web.cancel_compression = False
-            self.build_zipfile_list(filenames, processed_size_callback)
-
-        elif self.web.mode == 'website':
-            self.common.log("WebsiteModeWeb", "set_file_info")
-            self.web.cancel_compression = True
-
-        return True
-
+        self.set_file_info_custom(filenames, processed_size_callback)
 
     def build_file_list(self, filenames):
         """
@@ -163,56 +143,7 @@ class SendBaseModeWeb:
         return True
 
     def render_logic(self, path=''):
-        if path in self.files:
-            filesystem_path = self.files[path]
-
-            # If it's a directory
-            if os.path.isdir(filesystem_path):
-                # Is there an index.html?
-                index_path = os.path.join(path, 'index.html')
-                if self.web.mode == 'website' and index_path in self.files:
-                    # Render it
-                    dirname = os.path.dirname(self.files[index_path])
-                    basename = os.path.basename(self.files[index_path])
-
-                    return send_from_directory(dirname, basename)
-
-                else:
-                    # Otherwise, render directory listing
-                    filenames = []
-                    for filename in os.listdir(filesystem_path):
-                        if os.path.isdir(os.path.join(filesystem_path, filename)):
-                            filenames.append(filename + '/')
-                        else:
-                            filenames.append(filename)
-                    filenames.sort()
-                    return self.directory_listing(filenames, path, filesystem_path)
-
-            # If it's a file
-            elif os.path.isfile(filesystem_path):
-                dirname = os.path.dirname(filesystem_path)
-                basename = os.path.basename(filesystem_path)
-                return send_from_directory(dirname, basename)
-
-            # If it's not a directory or file, throw a 404
-            else:
-                return self.web.error404()
-        else:
-            # Special case loading /
-
-            if path == '':
-                index_path = 'index.html'
-                if self.web.mode == 'website' and index_path in self.files:
-                    # Render it
-                    dirname = os.path.dirname(self.files[index_path])
-                    basename = os.path.basename(self.files[index_path])
-                    return send_from_directory(dirname, basename)
-                else:
-                    # Root directory listing
-                    filenames = list(self.root_files)
-                    filenames.sort()
-                    return self.directory_listing(filenames, path)
-
-            else:
-                # If the path isn't found, throw a 404
-                return self.web.error404()
+        """
+        Inherited class will implement this.
+        """
+        pass
diff --git a/onionshare/web/share_mode.py b/onionshare/web/share_mode.py
index 1ed8c29a..8558a996 100644
--- a/onionshare/web/share_mode.py
+++ b/onionshare/web/share_mode.py
@@ -4,7 +4,7 @@ import tempfile
 import zipfile
 import mimetypes
 import gzip
-from flask import Response, request, render_template, make_response
+from flask import Response, request, render_template, make_response, send_from_directory
 
 from .send_base_mode import SendBaseModeWeb
 from .. import strings
@@ -14,11 +14,6 @@ class ShareModeWeb(SendBaseModeWeb):
     """
     All of the web logic for share mode
     """
-    def init(self):
-        self.common.log('ShareModeWeb', '__init__')
-
-        self.define_routes()
-
     def define_routes(self):
         """
         The web app routes for sharing files
@@ -47,7 +42,6 @@ class ShareModeWeb(SendBaseModeWeb):
 
             return self.render_logic(path)
 
-
         @self.web.app.route("/download")
         def download():
             """
@@ -171,6 +165,61 @@ class ShareModeWeb(SendBaseModeWeb):
                 r.headers.set('Content-Type', content_type)
             return r
 
+    def directory_listing_template(self, path, files, dirs):
+        return make_response(render_template(
+            'send.html',
+            file_info=self.file_info,
+            files=files,
+            dirs=dirs,
+            filename=os.path.basename(self.download_filename),
+            filesize=self.filesize,
+            filesize_human=self.common.human_readable_filesize(self.download_filesize),
+            is_zipped=self.is_zipped,
+            static_url_path=self.web.static_url_path))
+
+    def set_file_info_custom(self, filenames, processed_size_callback):
+        self.common.log("ShareModeWeb", "set_file_info_custom")
+        self.web.cancel_compression = False
+        self.build_zipfile_list(filenames, processed_size_callback)
+
+    def render_logic(self, path=''):
+        if path in self.files:
+            filesystem_path = self.files[path]
+
+            # If it's a directory
+            if os.path.isdir(filesystem_path):
+                # Render directory listing
+                filenames = []
+                for filename in os.listdir(filesystem_path):
+                    if os.path.isdir(os.path.join(filesystem_path, filename)):
+                        filenames.append(filename + '/')
+                    else:
+                        filenames.append(filename)
+                filenames.sort()
+                return self.directory_listing(filenames, path)
+
+            # If it's a file
+            elif os.path.isfile(filesystem_path):
+                dirname = os.path.dirname(filesystem_path)
+                basename = os.path.basename(filesystem_path)
+                return send_from_directory(dirname, basename)
+
+            # If it's not a directory or file, throw a 404
+            else:
+                return self.web.error404()
+        else:
+            # Special case loading /
+
+            if path == '':
+                # Root directory listing
+                filenames = list(self.root_files)
+                filenames.sort()
+                return self.directory_listing(filenames, path)
+
+            else:
+                # If the path isn't found, throw a 404
+                return self.web.error404()
+
     def build_zipfile_list(self, filenames, processed_size_callback=None):
         self.common.log("ShareModeWeb", "build_zipfile_list")
         for filename in filenames:
diff --git a/onionshare/web/website_mode.py b/onionshare/web/website_mode.py
index f38daa06..bb712a59 100644
--- a/onionshare/web/website_mode.py
+++ b/onionshare/web/website_mode.py
@@ -2,7 +2,7 @@ import os
 import sys
 import tempfile
 import mimetypes
-from flask import Response, request, render_template, make_response
+from flask import Response, request, render_template, make_response, send_from_directory
 
 from .send_base_mode import SendBaseModeWeb
 from .. import strings
@@ -12,16 +12,10 @@ class WebsiteModeWeb(SendBaseModeWeb):
     """
     All of the web logic for website mode
     """
-    def init(self):
-        self.common.log('WebsiteModeWeb', '__init__')
-        self.define_routes()
-
-
     def define_routes(self):
         """
         The web app routes for sharing a website
         """
-
         @self.web.app.route('/', defaults={'path': ''})
         @self.web.app.route('/<path:path>')
         def path_public(path):
@@ -43,3 +37,69 @@ class WebsiteModeWeb(SendBaseModeWeb):
             })
 
             return self.render_logic(path)
+
+    def directory_listing_template(self, path, files, dirs):
+        return make_response(render_template('listing.html',
+            path=path,
+            files=files,
+            dirs=dirs,
+            static_url_path=self.web.static_url_path))
+
+    def set_file_info_custom(self, filenames, processed_size_callback):
+        self.common.log("WebsiteModeWeb", "set_file_info_custom")
+        self.web.cancel_compression = True
+
+    def render_logic(self, path=''):
+        if path in self.files:
+            filesystem_path = self.files[path]
+
+            # If it's a directory
+            if os.path.isdir(filesystem_path):
+                # Is there an index.html?
+                index_path = os.path.join(path, 'index.html')
+                if index_path in self.files:
+                    # Render it
+                    dirname = os.path.dirname(self.files[index_path])
+                    basename = os.path.basename(self.files[index_path])
+
+                    return send_from_directory(dirname, basename)
+
+                else:
+                    # Otherwise, render directory listing
+                    filenames = []
+                    for filename in os.listdir(filesystem_path):
+                        if os.path.isdir(os.path.join(filesystem_path, filename)):
+                            filenames.append(filename + '/')
+                        else:
+                            filenames.append(filename)
+                    filenames.sort()
+                    return self.directory_listing(filenames, path, filesystem_path)
+
+            # If it's a file
+            elif os.path.isfile(filesystem_path):
+                dirname = os.path.dirname(filesystem_path)
+                basename = os.path.basename(filesystem_path)
+                return send_from_directory(dirname, basename)
+
+            # If it's not a directory or file, throw a 404
+            else:
+                return self.web.error404()
+        else:
+            # Special case loading /
+
+            if path == '':
+                index_path = 'index.html'
+                if index_path in self.files:
+                    # Render it
+                    dirname = os.path.dirname(self.files[index_path])
+                    basename = os.path.basename(self.files[index_path])
+                    return send_from_directory(dirname, basename)
+                else:
+                    # Root directory listing
+                    filenames = list(self.root_files)
+                    filenames.sort()
+                    return self.directory_listing(filenames, path, filesystem_path)
+
+            else:
+                # If the path isn't found, throw a 404
+                return self.web.error404()
diff --git a/onionshare_gui/mode/share_mode/threads.py b/onionshare_gui/mode/share_mode/threads.py
index 24e2c242..fed362eb 100644
--- a/onionshare_gui/mode/share_mode/threads.py
+++ b/onionshare_gui/mode/share_mode/threads.py
@@ -41,12 +41,8 @@ class CompressThread(QtCore.QThread):
         self.mode.common.log('CompressThread', 'run')
 
         try:
-            if self.mode.web.share_mode.set_file_info(self.mode.filenames, processed_size_callback=self.set_processed_size):
-                self.success.emit()
-            else:
-                # Cancelled
-                pass
-
+            self.mode.web.share_mode.set_file_info(self.mode.filenames, processed_size_callback=self.set_processed_size)
+            self.success.emit()
             self.mode.app.cleanup_filenames += self.mode.web.share_mode.cleanup_filenames
         except OSError as e:
             self.error.emit(e.strerror)
diff --git a/onionshare_gui/mode/website_mode/__init__.py b/onionshare_gui/mode/website_mode/__init__.py
index 50af4725..9f01cabc 100644
--- a/onionshare_gui/mode/website_mode/__init__.py
+++ b/onionshare_gui/mode/website_mode/__init__.py
@@ -165,12 +165,8 @@ class WebsiteMode(Mode):
         Step 3 in starting the server. Display large filesize
         warning, if applicable.
         """
-
-        if self.web.website_mode.set_file_info(self.filenames):
-            self.success.emit()
-        else:
-            # Cancelled
-            pass
+        self.web.website_mode.set_file_info(self.filenames)
+        self.success.emit()
 
     def start_server_error_custom(self):
         """