Use high level tempfile classes that clean up after themselves

This commit is contained in:
Micah Lee 2021-11-18 20:03:18 -08:00
parent 236bc33bc8
commit f5b60de2e6
No known key found for this signature in database
GPG Key ID: 403C2657CD994F73
2 changed files with 32 additions and 28 deletions

View File

@ -42,10 +42,11 @@ class SendBaseModeWeb:
self.is_zipped = False self.is_zipped = False
self.download_filename = None self.download_filename = None
self.download_filesize = None self.download_filesize = None
self.gzip_filename = None
self.gzip_filesize = None
self.zip_writer = 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 # If autostop_sharing, only allow one download at a time
self.download_in_progress = False self.download_in_progress = False
@ -192,12 +193,12 @@ class SendBaseModeWeb:
# gzip compress the individual file, if it hasn't already been compressed # gzip compress the individual file, if it hasn't already been compressed
if use_gzip: if use_gzip:
if filesystem_path not in self.gzip_individual_files: if filesystem_path not in self.gzip_individual_files:
gzip_filename = tempfile.mkstemp("wb+")[1] self.gzip_files.append(
self._gzip_compress(filesystem_path, gzip_filename, 6, None) tempfile.TemporaryFile("wb+", dir=self.common.build_tmp_dir())
self.gzip_individual_files[filesystem_path] = gzip_filename )
gzip_file = self.gzip_files[-1]
# Make sure the gzip file gets cleaned up when onionshare stops self._gzip_compress(filesystem_path, gzip_file.name, 6, None)
self.web.cleanup_filenames.append(gzip_filename) self.gzip_individual_files[filesystem_path] = gzip_file.name
file_to_download = self.gzip_individual_files[filesystem_path] file_to_download = self.gzip_individual_files[filesystem_path]
filesize = os.path.getsize(self.gzip_individual_files[filesystem_path]) filesize = os.path.getsize(self.gzip_individual_files[filesystem_path])

View File

@ -134,8 +134,12 @@ class ShareModeWeb(SendBaseModeWeb):
The web app routes for sharing files The web app routes for sharing files
""" """
@self.web.app.route("/", defaults={"path": ""}, methods=["GET"], provide_automatic_options=False) @self.web.app.route(
@self.web.app.route("/<path:path>", methods=["GET"], provide_automatic_options=False) "/", defaults={"path": ""}, methods=["GET"], provide_automatic_options=False
)
@self.web.app.route(
"/<path:path>", methods=["GET"], provide_automatic_options=False
)
def index(path): def index(path):
""" """
Render the template for the onionshare landing page. Render the template for the onionshare landing page.
@ -159,7 +163,9 @@ class ShareModeWeb(SendBaseModeWeb):
return self.render_logic(path) 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(): def download():
""" """
Download the zip file. Download the zip file.
@ -183,7 +189,7 @@ class ShareModeWeb(SendBaseModeWeb):
# and serve that # and serve that
use_gzip = self.should_use_gzip() use_gzip = self.should_use_gzip()
if use_gzip: if use_gzip:
file_to_download = self.gzip_filename file_to_download = self.gzip_file.name
self.filesize = self.gzip_filesize self.filesize = self.gzip_filesize
etag = self.gzip_etag etag = self.gzip_etag
else: else:
@ -286,7 +292,9 @@ class ShareModeWeb(SendBaseModeWeb):
if if_unmod: if if_unmod:
if_date = parse_date(if_unmod) if_date = parse_date(if_unmod)
if if_date and not if_date.tzinfo: 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: if if_date and if_date > last_modified:
abort(412) abort(412)
elif range_header is None: elif range_header is None:
@ -484,17 +492,16 @@ class ShareModeWeb(SendBaseModeWeb):
self.download_etag = make_etag(f) self.download_etag = make_etag(f)
# Compress the file with gzip now, so we don't have to do it on each request # 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_file = tempfile.TemporaryFile(
self._gzip_compress( "wb+", dir=self.common.build_tmp_dir()
self.download_filename, self.gzip_filename, 6, processed_size_callback
) )
self.gzip_filesize = os.path.getsize(self.gzip_filename) self._gzip_compress(
with open(self.gzip_filename, "rb") as f: 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) 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 self.is_zipped = False
else: else:
@ -519,10 +526,6 @@ class ShareModeWeb(SendBaseModeWeb):
with open(self.download_filename, "rb") as f: with open(self.download_filename, "rb") as f:
self.download_etag = make_etag(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 self.is_zipped = True
return True return True
@ -542,10 +545,10 @@ class ZipWriter(object):
if zip_filename: if zip_filename:
self.zip_filename = zip_filename self.zip_filename = zip_filename
else: else:
self.zip_temp_dir = tempfile.mkdtemp() self.zip_temp_dir = tempfile.TemporaryDirectory(
self.zip_filename = ( dir=self.common.build_tmp_dir()
f"{self.zip_temp_dir}/onionshare_{self.common.random_string(4, 6)}.zip"
) )
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.z = zipfile.ZipFile(self.zip_filename, "w", allowZip64=True)
self.processed_size_callback = processed_size_callback self.processed_size_callback = processed_size_callback