From 9663f781476c4c72fcf795088d9ad6836c6132f9 Mon Sep 17 00:00:00 2001 From: Micah Lee Date: Thu, 18 Nov 2021 20:03:18 -0800 Subject: [PATCH] Use high level tempfile classes that clean up after themselves --- cli/onionshare_cli/web/send_base_mode.py | 17 +++++----- cli/onionshare_cli/web/share_mode.py | 43 +++++++++++++----------- 2 files changed, 32 insertions(+), 28 deletions(-) diff --git a/cli/onionshare_cli/web/send_base_mode.py b/cli/onionshare_cli/web/send_base_mode.py index 27de598a..52720f6a 100644 --- a/cli/onionshare_cli/web/send_base_mode.py +++ b/cli/onionshare_cli/web/send_base_mode.py @@ -42,10 +42,11 @@ class SendBaseModeWeb: self.is_zipped = False self.download_filename = None self.download_filesize = None - self.gzip_filename = None - self.gzip_filesize = None self.zip_writer = None + # Store the tempfile objects here, so when they're garbage collected the files are deleted + self.gzip_files = [] + # If autostop_sharing, only allow one download at a time self.download_in_progress = False @@ -192,12 +193,12 @@ class SendBaseModeWeb: # gzip compress the individual file, if it hasn't already been compressed if use_gzip: if filesystem_path not in self.gzip_individual_files: - gzip_filename = tempfile.mkstemp("wb+")[1] - self._gzip_compress(filesystem_path, gzip_filename, 6, None) - self.gzip_individual_files[filesystem_path] = gzip_filename - - # Make sure the gzip file gets cleaned up when onionshare stops - self.web.cleanup_filenames.append(gzip_filename) + self.gzip_files.append( + tempfile.TemporaryFile("wb+", dir=self.common.build_tmp_dir()) + ) + gzip_file = self.gzip_files[-1] + self._gzip_compress(filesystem_path, gzip_file.name, 6, None) + self.gzip_individual_files[filesystem_path] = gzip_file.name 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 92a4c9af..09799679 100644 --- a/cli/onionshare_cli/web/share_mode.py +++ b/cli/onionshare_cli/web/share_mode.py @@ -134,8 +134,12 @@ class ShareModeWeb(SendBaseModeWeb): The web app routes for sharing files """ - @self.web.app.route("/", defaults={"path": ""}, methods=["GET"], provide_automatic_options=False) - @self.web.app.route("/", methods=["GET"], provide_automatic_options=False) + @self.web.app.route( + "/", defaults={"path": ""}, methods=["GET"], provide_automatic_options=False + ) + @self.web.app.route( + "/", methods=["GET"], provide_automatic_options=False + ) def index(path): """ Render the template for the onionshare landing page. @@ -159,7 +163,9 @@ class ShareModeWeb(SendBaseModeWeb): return self.render_logic(path) - @self.web.app.route("/download", methods=["GET"], provide_automatic_options=False) + @self.web.app.route( + "/download", methods=["GET"], provide_automatic_options=False + ) def download(): """ Download the zip file. @@ -183,7 +189,7 @@ class ShareModeWeb(SendBaseModeWeb): # and serve that use_gzip = self.should_use_gzip() if use_gzip: - file_to_download = self.gzip_filename + file_to_download = self.gzip_file.name self.filesize = self.gzip_filesize etag = self.gzip_etag else: @@ -286,7 +292,9 @@ class ShareModeWeb(SendBaseModeWeb): if if_unmod: if_date = parse_date(if_unmod) if if_date and not if_date.tzinfo: - if_date = if_date.replace(tzinfo=timezone.utc) # Compatible with Flask < 2.0.0 + if_date = if_date.replace( + tzinfo=timezone.utc + ) # Compatible with Flask < 2.0.0 if if_date and if_date > last_modified: abort(412) elif range_header is None: @@ -484,17 +492,16 @@ class ShareModeWeb(SendBaseModeWeb): self.download_etag = make_etag(f) # Compress the file with gzip now, so we don't have to do it on each request - self.gzip_filename = tempfile.mkstemp("wb+")[1] - self._gzip_compress( - self.download_filename, self.gzip_filename, 6, processed_size_callback + self.gzip_file = tempfile.TemporaryFile( + "wb+", dir=self.common.build_tmp_dir() ) - self.gzip_filesize = os.path.getsize(self.gzip_filename) - with open(self.gzip_filename, "rb") as f: + self._gzip_compress( + self.download_filename, self.gzip_file.name, 6, processed_size_callback + ) + self.gzip_filesize = os.path.getsize(self.gzip_file.name) + with open(self.gzip_file.name, "rb") as f: self.gzip_etag = make_etag(f) - # Make sure the gzip file gets cleaned up when onionshare stops - self.web.cleanup_filenames.append(self.gzip_filename) - self.is_zipped = False else: @@ -519,10 +526,6 @@ class ShareModeWeb(SendBaseModeWeb): with open(self.download_filename, "rb") as f: self.download_etag = make_etag(f) - # Make sure the zip file gets cleaned up when onionshare stops - 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 return True @@ -542,10 +545,10 @@ class ZipWriter(object): if zip_filename: self.zip_filename = zip_filename else: - self.zip_temp_dir = tempfile.mkdtemp() - self.zip_filename = ( - f"{self.zip_temp_dir}/onionshare_{self.common.random_string(4, 6)}.zip" + self.zip_temp_dir = tempfile.TemporaryDirectory( + dir=self.common.build_tmp_dir() ) + self.zip_filename = f"{self.zip_temp_dir.name}/onionshare_{self.common.random_string(4, 6)}.zip" self.z = zipfile.ZipFile(self.zip_filename, "w", allowZip64=True) self.processed_size_callback = processed_size_callback