Merge pull request #1344 from mig5/cleanup_web_files

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)

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()
 
 

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 = []

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])

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)

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 = []

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 == []

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()

desktop/src/onionshare/tab/mode/__init__.py

@@ -382,7 +382,7 @@ class Mode(QtWidgets.QWidget):
             except Exception:
                 # Probably we had no port to begin with (Onion service didn't start)
                 pass
-        self.app.cleanup()
+        self.web.cleanup()
 
         self.stop_server_custom()
 

desktop/src/onionshare/tab/mode/share_mode/threads.py

@@ -47,9 +47,6 @@ class CompressThread(QtCore.QThread):
                 self.mode.filenames, processed_size_callback=self.set_processed_size
             )
             self.success.emit()
-            self.mode.app.cleanup_filenames += (
-                self.mode.web.share_mode.cleanup_filenames
-            )
         except OSError as e:
             self.error.emit(e.strerror)
 

desktop/src/onionshare/tab/tab.py

@@ -668,7 +668,7 @@ class Tab(QtWidgets.QWidget):
         if self.close_dialog.clickedButton() == self.close_dialog.accept_button:
             self.common.log("Tab", "close_tab", "close, closing tab")
             self.get_mode().stop_server()
-            self.app.cleanup()
+            self.get_mode().web.cleanup()
             return True
         # Cancel
         else:
@@ -681,4 +681,4 @@ class Tab(QtWidgets.QWidget):
             self.get_mode().web.stop(self.get_mode().app.port)
             self.get_mode().web_thread.quit()
             self.get_mode().web_thread.wait()
-        self.app.cleanup()
+            self.get_mode().web.cleanup()