Merge branch 'develop' into 1275_fix_website_mode

6 files changed, 38 insertions, 33 deletions

diff --git a/cli/onionshare_cli/__init__.py b/cli/onionshare_cli/__init__.py
index 6a7a0fde..c046e472 100644
--- a/cli/onionshare_cli/__init__.py
+++ b/cli/onionshare_cli/__init__.py
@@ -442,7 +442,6 @@ def main(cwd=None):
         print("Compressing files.")
         try:
             web.share_mode.set_file_info(filenames)
-            app.cleanup_filenames += web.share_mode.cleanup_filenames
         except OSError as e:
             print(e.strerror)
             sys.exit(1)
@@ -536,7 +535,7 @@ def main(cwd=None):
         web.stop(app.port)
     finally:
         # Shutdown
-        app.cleanup()
+        web.cleanup()
         onion.cleanup()
 
 
diff --git a/cli/onionshare_cli/onionshare.py b/cli/onionshare_cli/onionshare.py
index e518d2fb..bd94100f 100644
--- a/cli/onionshare_cli/onionshare.py
+++ b/cli/onionshare_cli/onionshare.py
@@ -19,7 +19,6 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
 """
 
 import os
-import shutil
 from .common import AutoStopTimer
 
 
@@ -89,21 +88,3 @@ class OnionShare(object):
         Stop the onion service
         """
         self.onion.stop_onion_service(mode_settings)
-
-    def cleanup(self):
-        """
-        Shut everything down and clean up temporary files, etc.
-        """
-        self.common.log("OnionShare", "cleanup")
-
-        # Cleanup files
-        try:
-            for filename in self.cleanup_filenames:
-                if os.path.isfile(filename):
-                    os.remove(filename)
-                elif os.path.isdir(filename):
-                    shutil.rmtree(filename)
-        except Exception:
-            # Don't crash if file is still in use
-            pass
-        self.cleanup_filenames = []
diff --git a/cli/onionshare_cli/web/share_mode.py b/cli/onionshare_cli/web/share_mode.py
index 4dee0cee..95aec1ba 100644
--- a/cli/onionshare_cli/web/share_mode.py
+++ b/cli/onionshare_cli/web/share_mode.py
@@ -497,7 +497,7 @@ class ShareModeWeb(SendBaseModeWeb):
                 self.gzip_etag = make_etag(f)
 
             # Make sure the gzip file gets cleaned up when onionshare stops
-            self.cleanup_filenames.append(self.gzip_filename)
+            self.web.cleanup_filenames.append(self.gzip_filename)
 
             self.is_zipped = False
 
@@ -524,7 +524,8 @@ class ShareModeWeb(SendBaseModeWeb):
                 self.download_etag = make_etag(f)
 
             # Make sure the zip file gets cleaned up when onionshare stops
-            self.cleanup_filenames.append(self.zip_writer.zip_filename)
+            self.web.cleanup_filenames.append(self.zip_writer.zip_filename)
+            self.web.cleanup_filenames.append(self.zip_writer.zip_temp_dir)
 
             self.is_zipped = True
 
@@ -545,8 +546,9 @@ class ZipWriter(object):
         if zip_filename:
             self.zip_filename = zip_filename
         else:
+            self.zip_temp_dir = tempfile.mkdtemp()
             self.zip_filename = (
-                f"{tempfile.mkdtemp()}/onionshare_{self.common.random_string(4, 6)}.zip"
+                f"{self.zip_temp_dir}/onionshare_{self.common.random_string(4, 6)}.zip"
             )
 
         self.z = zipfile.ZipFile(self.zip_filename, "w", allowZip64=True)
diff --git a/cli/onionshare_cli/web/web.py b/cli/onionshare_cli/web/web.py
index da15c23b..d88a7e4e 100644
--- a/cli/onionshare_cli/web/web.py
+++ b/cli/onionshare_cli/web/web.py
@@ -21,6 +21,7 @@ import logging
 import os
 import queue
 import requests
+import shutil
 from distutils.version import LooseVersion as Version
 
 import flask
@@ -162,6 +163,8 @@ class Web:
             self.socketio.init_app(self.app)
             self.chat_mode = ChatModeWeb(self.common, self)
 
+        self.cleanup_filenames = []
+
     def get_mode(self):
         if self.mode == "share":
             return self.share_mode
@@ -423,3 +426,21 @@ class Web:
 
         # Reset any password that was in use
         self.password = None
+
+    def cleanup(self):
+        """
+        Shut everything down and clean up temporary files, etc.
+        """
+        self.common.log("Web", "cleanup")
+
+        # Cleanup files
+        try:
+            for filename in self.cleanup_filenames:
+                if os.path.isfile(filename):
+                    os.remove(filename)
+                elif os.path.isdir(filename):
+                    shutil.rmtree(filename)
+        except Exception:
+            # Don't crash if file is still in use
+            pass
+        self.cleanup_filenames = []
diff --git a/cli/tests/test_cli.py b/cli/tests/test_cli.py
index be71240b..d838f712 100644
--- a/cli/tests/test_cli.py
+++ b/cli/tests/test_cli.py
@@ -36,7 +36,6 @@ class TestOnionShare:
     def test_init(self, onionshare_obj):
         assert onionshare_obj.hidserv_dir is None
         assert onionshare_obj.onion_host is None
-        assert onionshare_obj.cleanup_filenames == []
         assert onionshare_obj.local_only is False
 
     def test_start_onion_service(self, onionshare_obj, mode_settings_obj):
@@ -48,11 +47,3 @@ class TestOnionShare:
         onionshare_obj.local_only = True
         onionshare_obj.start_onion_service("share", mode_settings_obj)
         assert onionshare_obj.onion_host == "127.0.0.1:{}".format(onionshare_obj.port)
-
-    def test_cleanup(self, onionshare_obj, temp_dir_1024, temp_file_1024):
-        onionshare_obj.cleanup_filenames = [temp_dir_1024, temp_file_1024]
-        onionshare_obj.cleanup()
-
-        assert os.path.exists(temp_dir_1024) is False
-        assert os.path.exists(temp_dir_1024) is False
-        assert onionshare_obj.cleanup_filenames == []
diff --git a/cli/tests/test_cli_web.py b/cli/tests/test_cli_web.py
index ec51654c..f8c96f9c 100644
--- a/cli/tests/test_cli_web.py
+++ b/cli/tests/test_cli_web.py
@@ -51,6 +51,7 @@ def web_obj(temp_dir, common_obj, mode, num_files=0):
     web.generate_password()
     web.running = True
 
+    web.cleanup_filenames == []
     web.app.testing = True
 
     # Share mode
@@ -345,6 +346,16 @@ class TestWeb:
             res.get_data()
             assert res.status_code == 200
 
+    def test_cleanup(self, common_obj, temp_dir_1024, temp_file_1024):
+        web = web_obj(temp_dir_1024, common_obj, "share", 3)
+
+        web.cleanup_filenames = [temp_dir_1024, temp_file_1024]
+        web.cleanup()
+
+        assert os.path.exists(temp_file_1024) is False
+        assert os.path.exists(temp_dir_1024) is False
+        assert web.cleanup_filenames == []
+
     def _make_auth_headers(self, password):
         auth = base64.b64encode(b"onionshare:" + password.encode()).decode()
         h = Headers()