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)
This commit is contained in:
Micah Lee 2021-05-05 02:35:14 +00:00 committed by GitHub
commit 058544d1e6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
10 changed files with 42 additions and 41 deletions

View File

@ -442,7 +442,6 @@ def main(cwd=None):
print("Compressing files.") print("Compressing files.")
try: try:
web.share_mode.set_file_info(filenames) web.share_mode.set_file_info(filenames)
app.cleanup_filenames += web.share_mode.cleanup_filenames
except OSError as e: except OSError as e:
print(e.strerror) print(e.strerror)
sys.exit(1) sys.exit(1)
@ -536,7 +535,7 @@ def main(cwd=None):
web.stop(app.port) web.stop(app.port)
finally: finally:
# Shutdown # Shutdown
app.cleanup() web.cleanup()
onion.cleanup() onion.cleanup()

View File

@ -19,7 +19,6 @@ along with this program. If not, see <http://www.gnu.org/licenses/>.
""" """
import os import os
import shutil
from .common import AutoStopTimer from .common import AutoStopTimer
@ -89,21 +88,3 @@ class OnionShare(object):
Stop the onion service Stop the onion service
""" """
self.onion.stop_onion_service(mode_settings) 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 = []

View File

@ -70,7 +70,6 @@ class SendBaseModeWeb:
self.root_files = ( self.root_files = (
{} {}
) # This is only the root files and dirs, as opposed to all of them ) # This is only the root files and dirs, as opposed to all of them
self.cleanup_filenames = []
self.cur_history_id = 0 self.cur_history_id = 0
self.file_info = {"files": [], "dirs": []} self.file_info = {"files": [], "dirs": []}
self.gzip_individual_files = {} self.gzip_individual_files = {}
@ -177,7 +176,7 @@ class SendBaseModeWeb:
self.gzip_individual_files[filesystem_path] = gzip_filename self.gzip_individual_files[filesystem_path] = gzip_filename
# Make sure the gzip file gets cleaned up when onionshare stops # 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] 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

@ -497,7 +497,7 @@ class ShareModeWeb(SendBaseModeWeb):
self.gzip_etag = make_etag(f) self.gzip_etag = make_etag(f)
# Make sure the gzip file gets cleaned up when onionshare stops # 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 self.is_zipped = False
@ -524,7 +524,8 @@ class ShareModeWeb(SendBaseModeWeb):
self.download_etag = make_etag(f) self.download_etag = make_etag(f)
# Make sure the zip file gets cleaned up when onionshare stops # 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 self.is_zipped = True
@ -545,8 +546,9 @@ 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_filename = ( 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) self.z = zipfile.ZipFile(self.zip_filename, "w", allowZip64=True)

View File

@ -21,6 +21,7 @@ import logging
import os import os
import queue import queue
import requests import requests
import shutil
from distutils.version import LooseVersion as Version from distutils.version import LooseVersion as Version
import flask import flask
@ -162,6 +163,8 @@ class Web:
self.socketio.init_app(self.app) self.socketio.init_app(self.app)
self.chat_mode = ChatModeWeb(self.common, self) self.chat_mode = ChatModeWeb(self.common, self)
self.cleanup_filenames = []
def get_mode(self): def get_mode(self):
if self.mode == "share": if self.mode == "share":
return self.share_mode return self.share_mode
@ -423,3 +426,21 @@ class Web:
# Reset any password that was in use # Reset any password that was in use
self.password = None 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 = []

View File

@ -36,7 +36,6 @@ class TestOnionShare:
def test_init(self, onionshare_obj): def test_init(self, onionshare_obj):
assert onionshare_obj.hidserv_dir is None assert onionshare_obj.hidserv_dir is None
assert onionshare_obj.onion_host is None assert onionshare_obj.onion_host is None
assert onionshare_obj.cleanup_filenames == []
assert onionshare_obj.local_only is False assert onionshare_obj.local_only is False
def test_start_onion_service(self, onionshare_obj, mode_settings_obj): def test_start_onion_service(self, onionshare_obj, mode_settings_obj):
@ -48,11 +47,3 @@ class TestOnionShare:
onionshare_obj.local_only = True onionshare_obj.local_only = True
onionshare_obj.start_onion_service("share", mode_settings_obj) onionshare_obj.start_onion_service("share", mode_settings_obj)
assert onionshare_obj.onion_host == "127.0.0.1:{}".format(onionshare_obj.port) 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 == []

View File

@ -51,6 +51,7 @@ def web_obj(temp_dir, common_obj, mode, num_files=0):
web.generate_password() web.generate_password()
web.running = True web.running = True
web.cleanup_filenames == []
web.app.testing = True web.app.testing = True
# Share mode # Share mode
@ -345,6 +346,16 @@ class TestWeb:
res.get_data() res.get_data()
assert res.status_code == 200 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): def _make_auth_headers(self, password):
auth = base64.b64encode(b"onionshare:" + password.encode()).decode() auth = base64.b64encode(b"onionshare:" + password.encode()).decode()
h = Headers() h = Headers()

View File

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

View File

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

View File

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