Start code sharing between WebsiteMode and ShareMode

3 files changed, 52 insertions, 36 deletions

diff --git a/onionshare/web/base_mode.py b/onionshare/web/base_mode.py
new file mode 100644
index 00000000..fb1043d7
--- /dev/null
+++ b/onionshare/web/base_mode.py
@@ -0,0 +1,45 @@
+import os
+import sys
+import tempfile
+import mimetypes
+from flask import Response, request, render_template, make_response
+
+from .. import strings
+
+class BaseModeWeb(object):
+    """
+    All of the web logic shared between share and website mode
+    """
+    def __init__(self, common, web):
+        super(BaseModeWeb, self).__init__()
+        self.common = common
+        self.web = web
+
+        # Information about the file to be shared
+        self.file_info = []
+        self.is_zipped = False
+        self.download_filename = None
+        self.download_filesize = None
+        self.gzip_filename = None
+        self.gzip_filesize = None
+        self.zip_writer = None
+
+        # Dictionary mapping file paths to filenames on disk
+        self.files = {}
+        
+        self.visit_count = 0
+        self.download_count = 0
+
+        # If "Stop After First Download" is checked (stay_open == False), only allow
+        # one download at a time.
+        self.download_in_progress = False
+
+        # Reset assets path
+        self.web.app.static_folder=self.common.get_resource_path('static')
+
+
+    def init(self):
+        """
+        Add custom initialization here.
+        """
+        pass
diff --git a/onionshare/web/share_mode.py b/onionshare/web/share_mode.py
index 0dfa7e0a..779d0a4b 100644
--- a/onionshare/web/share_mode.py
+++ b/onionshare/web/share_mode.py
@@ -6,38 +6,17 @@ import mimetypes
 import gzip
 from flask import Response, request, render_template, make_response
 
+from .base_mode import BaseModeWeb
 from .. import strings
 
 
-class ShareModeWeb(object):
+class ShareModeWeb(BaseModeWeb):
     """
     All of the web logic for share mode
     """
-    def __init__(self, common, web):
-        self.common = common
+    def init(self):
         self.common.log('ShareModeWeb', '__init__')
 
-        self.web = web
-
-        # Information about the file to be shared
-        self.file_info = []
-        self.is_zipped = False
-        self.download_filename = None
-        self.download_filesize = None
-        self.gzip_filename = None
-        self.gzip_filesize = None
-        self.zip_writer = None
-
-        self.download_count = 0
-
-        # If "Stop After First Download" is checked (stay_open == False), only allow
-        # 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 define_routes(self):
diff --git a/onionshare/web/website_mode.py b/onionshare/web/website_mode.py
index d2cd6db9..f61da569 100644
--- a/onionshare/web/website_mode.py
+++ b/onionshare/web/website_mode.py
@@ -4,25 +4,17 @@ import tempfile
 import mimetypes
 from flask import Response, request, render_template, make_response, send_from_directory
 
+from .base_mode import BaseModeWeb
 from .. import strings
 
 
-class WebsiteModeWeb(object):
+class WebsiteModeWeb(BaseModeWeb):
     """
     All of the web logic for share mode
     """
-    def __init__(self, common, web):
-        self.common = common
-        self.common.log('WebsiteModeWeb', '__init__')
-
-        self.web = web
+    def init(self):
 
-        # Dictionary mapping file paths to filenames on disk
-        self.files = {}
-        self.visit_count = 0
-
-        # Reset assets path
-        self.web.app.static_folder=self.common.get_resource_path('static')
+        self.common.log('WebsiteModeWeb', '__init__')
 
         self.define_routes()