From d9e48d6f67f0d2aa855781ba77837cbb9929676a Mon Sep 17 00:00:00 2001 From: Kunal Mehta Date: Tue, 26 Nov 2024 12:39:22 -0800 Subject: [PATCH] WIP: Display download progress for file downloads --- .../securedrop_client/api_jobs/downloads.py | 9 ++++--- client/securedrop_client/gui/widgets.py | 27 +++++++++++-------- client/securedrop_client/logic.py | 15 ++++++++--- client/securedrop_client/sdk/__init__.py | 26 +++++++++++++++--- client/securedrop_client/sdk/progress.py | 10 +++++++ client/tests/api_jobs/test_downloads.py | 16 ++++++----- client/tests/gui/test_widgets.py | 25 ----------------- 7 files changed, 73 insertions(+), 55 deletions(-) create mode 100644 client/securedrop_client/sdk/progress.py diff --git a/client/securedrop_client/api_jobs/downloads.py b/client/securedrop_client/api_jobs/downloads.py index 7499bd69b..5e6af6948 100644 --- a/client/securedrop_client/api_jobs/downloads.py +++ b/client/securedrop_client/api_jobs/downloads.py @@ -12,7 +12,7 @@ from securedrop_client.api_jobs.base import SingleObjectApiJob from securedrop_client.crypto import CryptoError, GpgHelper from securedrop_client.db import DownloadError, DownloadErrorCodes, File, Message, Reply -from securedrop_client.sdk import API, BaseError +from securedrop_client.sdk import API, BaseError, Progress from securedrop_client.sdk import Reply as SdkReply from securedrop_client.sdk import Submission as SdkSubmission from securedrop_client.storage import ( @@ -56,6 +56,7 @@ class DownloadJob(SingleObjectApiJob): def __init__(self, data_dir: str, uuid: str) -> None: super().__init__(uuid) self.data_dir = data_dir + self.progress: Progress | None = None def _get_realistic_timeout(self, size_in_bytes: int) -> int: """ @@ -260,7 +261,7 @@ def call_download_api(self, api: API, db_object: Reply) -> tuple[str, str]: # will want to pass the default request timeout to download_reply instead of setting it on # the api object directly. api.default_request_timeout = 20 - return api.download_reply(sdk_object) + return api.download_reply(sdk_object, progress=self.progress) def call_decrypt(self, filepath: str, session: Session | None = None) -> str: """ @@ -316,7 +317,7 @@ def call_download_api(self, api: API, db_object: Message) -> tuple[str, str]: sdk_object.source_uuid = db_object.source.uuid sdk_object.filename = db_object.filename return api.download_submission( - sdk_object, timeout=self._get_realistic_timeout(db_object.size) + sdk_object, timeout=self._get_realistic_timeout(db_object.size), progress=self.progress ) def call_decrypt(self, filepath: str, session: Session | None = None) -> str: @@ -375,7 +376,7 @@ def call_download_api(self, api: API, db_object: File) -> tuple[str, str]: sdk_object.source_uuid = db_object.source.uuid sdk_object.filename = db_object.filename return api.download_submission( - sdk_object, timeout=self._get_realistic_timeout(db_object.size) + sdk_object, timeout=self._get_realistic_timeout(db_object.size), progress=self.progress ) def call_decrypt(self, filepath: str, session: Session | None = None) -> str: diff --git a/client/securedrop_client/gui/widgets.py b/client/securedrop_client/gui/widgets.py index a92955f4f..b7d14864c 100644 --- a/client/securedrop_client/gui/widgets.py +++ b/client/securedrop_client/gui/widgets.py @@ -51,6 +51,7 @@ QListWidgetItem, QMenu, QPlainTextEdit, + QProgressBar, QPushButton, QScrollArea, QSizePolicy, @@ -90,6 +91,7 @@ from securedrop_client.gui.source import DeleteSourceDialog from securedrop_client.logic import Controller from securedrop_client.resources import load_css, load_icon, load_image, load_movie +from securedrop_client.sdk import Progress from securedrop_client.storage import source_exists from securedrop_client.utils import humanize_filesize @@ -2579,6 +2581,9 @@ def __init__( self.download_button.setIcon(load_icon("download_file.svg")) self.download_button.setFont(self.file_buttons_font) self.download_button.setCursor(QCursor(Qt.PointingHandCursor)) + self.download_progress = QProgressBar() + self.download_progress.setMaximum(self.file.size) + self.download_progress.hide() self.download_animation = load_movie("download_file.gif") self.export_button = QPushButton(_("EXPORT")) self.export_button.setObjectName("FileWidget_export_print") @@ -2590,6 +2595,7 @@ def __init__( self.print_button.setFont(self.file_buttons_font) self.print_button.setCursor(QCursor(Qt.PointingHandCursor)) file_options_layout.addWidget(self.download_button) + file_options_layout.addWidget(self.download_progress) file_options_layout.addWidget(self.export_button) file_options_layout.addWidget(self.middot) file_options_layout.addWidget(self.print_button) @@ -2675,6 +2681,7 @@ def _set_file_state(self) -> None: logger.debug(f"Changing file {self.uuid} state to decrypted/downloaded") self._set_file_name() self.download_button.hide() + self.download_progress.hide() self.no_file_name.hide() self.export_button.show() self.middot.show() @@ -2693,6 +2700,7 @@ def _set_file_state(self) -> None: self.download_button.setFont(self.file_buttons_font) self.download_button.show() + self.download_progress.hide() # Reset stylesheet self.download_button.setStyleSheet("") @@ -2793,15 +2801,17 @@ def _on_left_click(self) -> None: if self.controller.api: self.start_button_animation() # Download the file. - self.controller.on_submission_download(File, self.uuid) + self.controller.on_submission_download( + File, self.uuid, Progress(self.download_progress) + ) def start_button_animation(self) -> None: """ Update the download button to the animated "downloading" state. """ self.downloading = True - self.download_animation.frameChanged.connect(self.set_button_animation_frame) - self.download_animation.start() + self.download_progress.setValue(0) + self.download_progress.show() self.download_button.setText(_(" DOWNLOADING ")) # Reset widget stylesheet @@ -2809,18 +2819,13 @@ def start_button_animation(self) -> None: self.download_button.setObjectName("FileWidget_download_button_animating") self.download_button.setStyleSheet(self.DOWNLOAD_BUTTON_CSS) - def set_button_animation_frame(self, frame_number: int) -> None: - """ - Sets the download button's icon to the current frame of the spinner - animation. - """ - self.download_button.setIcon(QIcon(self.download_animation.currentPixmap())) - def stop_button_animation(self) -> None: """ Stops the download animation and restores the button to its default state. """ - self.download_animation.stop() + self.download_progress.hide() + # FIXME: will this stop the segfaults?? + self.download_progress.deleteLater() file = self.controller.get_file(self.file.uuid) if file is None: self.deleteLater() diff --git a/client/securedrop_client/logic.py b/client/securedrop_client/logic.py index dc5b04275..cd77bcb87 100644 --- a/client/securedrop_client/logic.py +++ b/client/securedrop_client/logic.py @@ -60,7 +60,7 @@ ) from securedrop_client.crypto import GpgHelper from securedrop_client.queue import ApiJobQueue -from securedrop_client.sdk import AuthError, RequestTimeoutError, ServerConnectionError +from securedrop_client.sdk import AuthError, Progress, RequestTimeoutError, ServerConnectionError from securedrop_client.sync import ApiSync from securedrop_client.utils import check_dir_permissions @@ -825,7 +825,10 @@ def set_status(self, message: str, duration: int = 5000) -> None: @login_required def _submit_download_job( - self, object_type: type[db.Reply] | type[db.Message] | type[db.File], uuid: str + self, + object_type: type[db.Reply] | type[db.Message] | type[db.File], + uuid: str, + progress: Progress | None = None, ) -> None: if object_type == db.Reply: job: ReplyDownloadJob | MessageDownloadJob | FileDownloadJob = ReplyDownloadJob( @@ -842,6 +845,7 @@ def _submit_download_job( job.success_signal.connect(self.on_file_download_success) job.failure_signal.connect(self.on_file_download_failure) + job.progress = progress self.add_job.emit(job) def download_new_messages(self) -> None: @@ -974,12 +978,15 @@ def on_file_open(self, file: db.File) -> None: @login_required def on_submission_download( - self, submission_type: type[db.File] | type[db.Message], submission_uuid: str + self, + submission_type: type[db.File] | type[db.Message], + submission_uuid: str, + progress: Progress | None = None, ) -> None: """ Download the file associated with the Submission (which may be a File or Message). """ - self._submit_download_job(submission_type, submission_uuid) + self._submit_download_job(submission_type, submission_uuid, progress) def on_file_download_success(self, uuid: str) -> None: """ diff --git a/client/securedrop_client/sdk/__init__.py b/client/securedrop_client/sdk/__init__.py index 810e3d9a0..76dc604ba 100644 --- a/client/securedrop_client/sdk/__init__.py +++ b/client/securedrop_client/sdk/__init__.py @@ -11,6 +11,7 @@ from securedrop_client.config import Config +from .progress import Progress from .sdlocalobjects import ( AuthError, BaseError, @@ -131,7 +132,10 @@ def _rpc_target(self) -> list: return [f"{target_directory}/debug/securedrop-proxy"] def _streaming_download( - self, data: dict[str, Any], env: dict + self, + data: dict[str, Any], + env: dict, + progress: Progress, ) -> StreamedResponse | JSONResponse: fobj = tempfile.TemporaryFile("w+b") @@ -183,6 +187,7 @@ def _streaming_download( fobj.write(chunk) bytes_written += len(chunk) + progress.set_value(bytes_written) logger.debug(f"Retry {retry}, bytes written: {bytes_written:,}") # Wait for the process to end @@ -303,6 +308,7 @@ def _send_json_request( body: str | None = None, headers: dict[str, str] | None = None, timeout: int | None = None, + progress: Progress | None = None, ) -> StreamedResponse | JSONResponse: """Build a JSON-serialized request to pass to the proxy. Handle the JSON or streamed response back, plus translate HTTP error statuses @@ -328,9 +334,12 @@ def _send_json_request( if self.development_mode: env["SD_PROXY_ORIGIN"] = self.server + if not progress: + progress = Progress(None) + # Streaming if stream: - return self._streaming_download(data, env) + return self._streaming_download(data, env, progress) # Not streaming data_str = json.dumps(data).encode() @@ -360,6 +369,7 @@ def _send_json_request( logger.error("Internal proxy error (non-streaming)") raise BaseError(f"Internal proxy error: {error_desc}") + progress.set_value(len(response.stdout)) return self._handle_json_response(response.stdout) def authenticate(self, totp: str | None = None) -> bool: @@ -683,7 +693,11 @@ def delete_submission(self, submission: Submission) -> bool: return False def download_submission( - self, submission: Submission, path: str | None = None, timeout: int | None = None + self, + submission: Submission, + path: str | None = None, + timeout: int | None = None, + progress: Progress | None = None, ) -> tuple[str, str]: """ Returns a tuple of etag (format is algorithm:checksum) and file path for @@ -712,6 +726,7 @@ def download_submission( stream=True, headers=self.build_headers(), timeout=timeout or self.default_download_timeout, + progress=progress, ) if isinstance(response, JSONResponse): @@ -909,7 +924,9 @@ def get_all_replies(self) -> list[Reply]: return result - def download_reply(self, reply: Reply, path: str | None = None) -> tuple[str, str]: + def download_reply( + self, reply: Reply, path: str | None = None, progress: Progress | None = None + ) -> tuple[str, str]: """ Returns a tuple of etag (format is algorithm:checksum) and file path for a given Reply object. This method requires a directory path @@ -936,6 +953,7 @@ def download_reply(self, reply: Reply, path: str | None = None) -> tuple[str, st stream=True, headers=self.build_headers(), timeout=self.default_request_timeout, + progress=progress, ) if isinstance(response, JSONResponse): diff --git a/client/securedrop_client/sdk/progress.py b/client/securedrop_client/sdk/progress.py new file mode 100644 index 000000000..c6a7d1309 --- /dev/null +++ b/client/securedrop_client/sdk/progress.py @@ -0,0 +1,10 @@ +from PyQt5.QtWidgets import QProgressBar + + +class Progress: + def __init__(self, inner: QProgressBar | None) -> None: + self.inner = inner + + def set_value(self, value: int) -> None: + if self.inner: + self.inner.setValue(value) diff --git a/client/tests/api_jobs/test_downloads.py b/client/tests/api_jobs/test_downloads.py index bd1c5cf77..6bc0cde5f 100644 --- a/client/tests/api_jobs/test_downloads.py +++ b/client/tests/api_jobs/test_downloads.py @@ -12,7 +12,7 @@ ReplyDownloadJob, ) from securedrop_client.crypto import CryptoError, GpgHelper -from securedrop_client.sdk import BaseError +from securedrop_client.sdk import BaseError, Progress from securedrop_client.sdk import Submission as SdkSubmission from tests import factory @@ -351,7 +351,9 @@ def test_FileDownloadJob_happy_path_no_etag(mocker, homedir, session, session_ma gpg = GpgHelper(homedir, session_maker, is_qubes=False) mock_decrypt = patch_decrypt(mocker, homedir, gpg, file_.filename) - def fake_download(sdk_obj: SdkSubmission, timeout: int) -> tuple[str, str]: + def fake_download( + sdk_obj: SdkSubmission, timeout: int, progress: Progress | None + ) -> tuple[str, str]: """ :return: (etag, path_to_dl) """ @@ -389,7 +391,7 @@ def test_FileDownloadJob_happy_path_sha256_etag( gpg = GpgHelper(homedir, session_maker, is_qubes=False) mock_decrypt = patch_decrypt(mocker, homedir, gpg, file_.filename) - def fake_download(sdk_obj: SdkSubmission, timeout: int) -> tuple[str, str]: + def fake_download(sdk_obj: SdkSubmission, timeout: int, progress: Progress) -> tuple[str, str]: """ :return: (etag, path_to_dl) """ @@ -426,7 +428,7 @@ def test_FileDownloadJob_bad_sha256_etag( gpg = GpgHelper(homedir, session_maker, is_qubes=False) - def fake_download(sdk_obj: SdkSubmission, timeout: int) -> tuple[str, str]: + def fake_download(sdk_obj: SdkSubmission, timeout: int, progress: Progress) -> tuple[str, str]: """ :return: (etag, path_to_dl) """ @@ -455,7 +457,7 @@ def test_FileDownloadJob_happy_path_unknown_etag(mocker, homedir, session, sessi gpg = GpgHelper(homedir, session_maker, is_qubes=False) - def fake_download(sdk_obj: SdkSubmission, timeout: int) -> tuple[str, str]: + def fake_download(sdk_obj: SdkSubmission, timeout: int, progress: Progress) -> tuple[str, str]: """ :return: (etag, path_to_dl) """ @@ -494,7 +496,7 @@ def test_FileDownloadJob_decryption_error( gpg = GpgHelper(homedir, session_maker, is_qubes=False) mock_decrypt = mocker.patch.object(gpg, "decrypt_submission_or_reply", side_effect=CryptoError) - def fake_download(sdk_obj: SdkSubmission, timeout: int) -> tuple[str, str]: + def fake_download(sdk_obj: SdkSubmission, timeout: int, progress: Progress) -> tuple[str, str]: """ :return: (etag, path_to_dl) """ @@ -535,7 +537,7 @@ def test_FileDownloadJob_raises_on_path_traversal_attack( api_client = mocker.MagicMock() download_fn = mocker.patch.object(api_client, "download_reply") - def fake_download(sdk_obj: SdkSubmission, timeout: int) -> tuple[str, str]: + def fake_download(sdk_obj: SdkSubmission, timeout: int, progress: Progress) -> tuple[str, str]: """ :return: (etag, path-to-download) """ diff --git a/client/tests/gui/test_widgets.py b/client/tests/gui/test_widgets.py index 8bf90a623..b0b13aa80 100644 --- a/client/tests/gui/test_widgets.py +++ b/client/tests/gui/test_widgets.py @@ -3416,31 +3416,6 @@ def test_FileWidget_on_left_click_open(mocker, session, source): fw.controller.on_file_open.assert_called_once_with(file_) -def test_FileWidget_set_button_animation_frame(mocker, session, source): - """ - Left click on download when file is not downloaded should trigger - a download. - """ - file_ = factory.File(source=source["source"], is_downloaded=False, is_decrypted=None) - session.add(file_) - session.commit() - - controller = mocker.MagicMock() - - fw = FileWidget( - file_, - controller, - mocker.MagicMock(), - mocker.MagicMock(), - mocker.MagicMock(), - 0, - 123, - ) - fw.download_button = mocker.MagicMock() - fw.set_button_animation_frame(1) - assert fw.download_button.setIcon.call_count == 1 - - def test_FileWidget_update(mocker, session, source): """ The update method should show/hide widgets if file is downloaded