Refactor directory_listing function

3 files changed, 50 insertions, 41 deletions

diff --git a/onionshare/web/base_mode.py b/onionshare/web/base_mode.py
index 46e63e1a..ff7e11be 100644
--- a/onionshare/web/base_mode.py
+++ b/onionshare/web/base_mode.py
@@ -4,7 +4,7 @@ import tempfile
 import mimetypes
 from flask import Response, request, render_template, make_response
 
- from .. import strings
+from .. import strings
 
 class BaseModeWeb(object):
     """
@@ -46,36 +46,31 @@ class BaseModeWeb(object):
         pass
 
 
-    def directory_listing(self, path, filenames, filesystem_path=None):
+    def directory_listing(self, path='', filenames=[], filesystem_path=None):
         # If filesystem_path is None, this is the root directory listing
         files = []
         dirs = []
+        r = ''
+
+        if self.web.mode == 'website':
+            files, dirs = build_directory_listing(filenames)
+
+            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,
+                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))
 
-        for filename in filenames:
-            if filesystem_path:
-                this_filesystem_path = os.path.join(filesystem_path, filename)
-            else:
-                this_filesystem_path = self.files[filename]
-
-            is_dir = os.path.isdir(this_filesystem_path)
-
-            if is_dir:
-                dirs.append({
-                    'basename': filename
-                })
-            else:
-                size = os.path.getsize(this_filesystem_path)
-                size_human = self.common.human_readable_filesize(size)
-                files.append({
-                    'basename': filename,
-                    'size_human': size_human
-                })
-
-        r = make_response(render_template('listing.html',
-            path=path,
-            files=files,
-            dirs=dirs,
-            static_url_path=self.web.static_url_path))
         return self.web.add_security_headers(r)
 
 
diff --git a/onionshare/web/share_mode.py b/onionshare/web/share_mode.py
index 68763357..cb3bba50 100644
--- a/onionshare/web/share_mode.py
+++ b/onionshare/web/share_mode.py
@@ -44,15 +44,8 @@ class ShareModeWeb(BaseModeWeb):
             else:
                 self.filesize = self.download_filesize
 
-            r = make_response(render_template(
-                'send.html',
-                file_info=self.file_info,
-                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))
-            return self.web.add_security_headers(r)
+            return self.directory_listing()
+
 
         @self.web.app.route("/download")
         def download():
diff --git a/onionshare/web/website_mode.py b/onionshare/web/website_mode.py
index 287acbd9..5183bebc 100644
--- a/onionshare/web/website_mode.py
+++ b/onionshare/web/website_mode.py
@@ -14,12 +14,9 @@ class WebsiteModeWeb(BaseModeWeb):
     """
     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):
         """
         The web app routes for sharing a website
@@ -56,6 +53,7 @@ class WebsiteModeWeb(BaseModeWeb):
                         # 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:
@@ -80,6 +78,7 @@ class WebsiteModeWeb(BaseModeWeb):
                     return self.web.error404()
             else:
                 # Special case loading /
+
                 if path == '':
                     index_path = 'index.html'
                     if index_path in self.files:
@@ -97,7 +96,29 @@ class WebsiteModeWeb(BaseModeWeb):
                     # If the path isn't found, throw a 404
                     return self.web.error404()
 
-    
+    def build_directory_listing(self, filenames):
+        for filename in filenames:
+            if filesystem_path:
+                this_filesystem_path = os.path.join(filesystem_path, filename)
+            else:
+                this_filesystem_path = self.files[filename]
+
+            is_dir = os.path.isdir(this_filesystem_path)
+
+            if is_dir:
+                dirs.append({
+                    'basename': filename
+                })
+            else:
+                size = os.path.getsize(this_filesystem_path)
+                size_human = self.common.human_readable_filesize(size)
+                files.append({
+                    'basename': filename,
+                    'size_human': size_human
+                })
+        return files, dirs
+
+
     def build_file_list(self, filenames):
         """
         Build a data structure that describes the list of files that make up