Move the cleanup() function from Onionshare class to Web class, so that the list of files to be cleaned up is always available (needed for website temp files)

7 files changed, 36 insertions, 34 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/send_base_mode.py b/cli/onionshare_cli/web/send_base_mode.py
index c7347347..77cb8ba5 100644
--- a/cli/onionshare_cli/web/send_base_mode.py
+++ b/cli/onionshare_cli/web/send_base_mode.py
@@ -70,7 +70,6 @@ class SendBaseModeWeb:
         self.root_files = (
             {}
         )  # This is only the root files and dirs, as opposed to all of them
-        self.cleanup_filenames = []
         self.cur_history_id = 0
         self.file_info = {"files": [], "dirs": []}
         self.gzip_individual_files = {}
@@ -177,7 +176,7 @@ class SendBaseModeWeb:
                 self.gzip_individual_files[filesystem_path] = gzip_filename
 
                 # Make sure the gzip file gets cleaned up when onionshare stops
-                self.cleanup_filenames.append(gzip_filename)
+                self.web.cleanup_filenames.append(gzip_filename)
 
             file_to_download = self.gzip_individual_files[filesystem_path]
             filesize = os.path.getsize(self.gzip_individual_files[filesystem_path])
diff --git a/cli/onionshare_cli/web/share_mode.py b/cli/onionshare_cli/web/share_mode.py
index 4dee0cee..c5007d7f 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,7 @@ 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.is_zipped = 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()