Refactor set_file_list between website and share mode

3 files changed, 36 insertions, 30 deletions

diff --git a/onionshare/web/base_mode.py b/onionshare/web/base_mode.py
index fb1043d7..8843d198 100644
--- a/onionshare/web/base_mode.py
+++ b/onionshare/web/base_mode.py
@@ -26,7 +26,9 @@ class BaseModeWeb(object):
 
         # Dictionary mapping file paths to filenames on disk
         self.files = {}
-        
+        self.cleanup_filenames = []
+        self.file_info = {'files': [], 'dirs': []}
+
         self.visit_count = 0
         self.download_count = 0
 
@@ -34,8 +36,7 @@ class BaseModeWeb(object):
         # one download at a time.
         self.download_in_progress = False
 
-        # Reset assets path
-        self.web.app.static_folder=self.common.get_resource_path('static')
+        self.define_routes()
 
 
     def init(self):
@@ -43,3 +44,27 @@ class BaseModeWeb(object):
         Add custom initialization here.
         """
         pass
+
+    def set_file_info(self, filenames, processed_size_callback=None):
+        """
+        Build a data structure that describes the list of files
+        """
+        if self.web.mode == 'website':
+            self.common.log("WebsiteModeWeb", "set_file_info")
+            self.web.cancel_compression = True
+
+            # This is only the root files and dirs, as opposed to all of them
+            self.root_files = {}
+
+            # If there's just one folder, replace filenames with a list of files inside that folder
+            if len(filenames) == 1 and os.path.isdir(filenames[0]):
+                filenames = [os.path.join(filenames[0], x) for x in os.listdir(filenames[0])]
+
+            self.build_file_list(filenames)
+
+        elif self.web.mode == 'share':
+            self.common.log("ShareModeWeb", "set_file_info")
+            self.web.cancel_compression = False
+            self.build_zipfile_list(filenames, processed_size_callback)
+
+        return True
diff --git a/onionshare/web/share_mode.py b/onionshare/web/share_mode.py
index 779d0a4b..68763357 100644
--- a/onionshare/web/share_mode.py
+++ b/onionshare/web/share_mode.py
@@ -177,19 +177,8 @@ class ShareModeWeb(BaseModeWeb):
                 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.common.log("ShareModeWeb", "set_file_info")
-        self.web.cancel_compression = False
-
-        self.cleanup_filenames = []
-
-        # build file info list
-        self.file_info = {'files': [], 'dirs': []}
+    def build_zipfile_list(self, filenames, processed_size_callback=None):
+        self.common.log("ShareModeWeb", "build_file_list")
         for filename in filenames:
             info = {
                 'filename': filename,
diff --git a/onionshare/web/website_mode.py b/onionshare/web/website_mode.py
index f61da569..4c024849 100644
--- a/onionshare/web/website_mode.py
+++ b/onionshare/web/website_mode.py
@@ -10,12 +10,14 @@ from .. import strings
 
 class WebsiteModeWeb(BaseModeWeb):
     """
-    All of the web logic for share mode
+    All of the web logic for website mode
     """
     def init(self):
-
         self.common.log('WebsiteModeWeb', '__init__')
 
+        # Reset assets path
+        self.web.app.static_folder=self.common.get_resource_path('share/static')
+
         self.define_routes()
 
     def define_routes(self):
@@ -127,22 +129,12 @@ class WebsiteModeWeb(BaseModeWeb):
             static_url_path=self.web.static_url_path))
         return self.web.add_security_headers(r)
 
-    def set_file_info(self, filenames):
+    def build_file_list(self, filenames):
         """
         Build a data structure that describes the list of files that make up
         the static website.
         """
-        self.common.log("WebsiteModeWeb", "set_file_info")
-
-        # This is a dictionary that maps HTTP routes to filenames on disk
-        self.files = {}
-
-        # This is only the root files and dirs, as opposed to all of them
-        self.root_files = {}
-
-        # If there's just one folder, replace filenames with a list of files inside that folder
-        if len(filenames) == 1 and os.path.isdir(filenames[0]):
-            filenames = [os.path.join(filenames[0], x) for x in os.listdir(filenames[0])]
+        self.common.log("WebsiteModeWeb", "build_file_list")
 
         # Loop through the files
         for filename in filenames: