diff --git a/selfprivacy_api/actions/system.py b/selfprivacy_api/actions/system.py index 853662f..f5e0dc0 100644 --- a/selfprivacy_api/actions/system.py +++ b/selfprivacy_api/actions/system.py @@ -2,7 +2,7 @@ import os import subprocess import pytz -from typing import Optional +from typing import Optional, List from pydantic import BaseModel from selfprivacy_api.utils import WriteUserData, ReadUserData @@ -58,36 +58,56 @@ def set_auto_upgrade_settings( user_data["autoUpgrade"]["allowReboot"] = allowReboot +class ShellException(Exception): + """Something went wrong when calling another process""" + + pass + + +def run_blocking(cmd: List[str], new_session: bool = False) -> str: + """Run a process, block until done, return output, complain if failed""" + process_handle = subprocess.Popen( + cmd, + shell=False, + start_new_session=new_session, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + stdout_raw, stderr_raw = process_handle.communicate() + stdout = stdout_raw.decode("utf-8") + if stderr_raw is not None: + stderr = stderr_raw.decode("utf-8") + else: + stderr = "" + output = stdout + "\n" + stderr + if process_handle.returncode != 0: + raise ShellException( + f"Shell command failed, command array: {cmd}, output: {output}" + ) + return stdout + + def rebuild_system() -> int: """Rebuild the system""" - rebuild_result = subprocess.Popen( - ["systemctl", "start", "sp-nixos-rebuild.service"], start_new_session=True - ) - rebuild_result.communicate()[0] - return rebuild_result.returncode + run_blocking(["systemctl", "start", "sp-nixos-rebuild.service"], new_session=True) + return 0 def rollback_system() -> int: """Rollback the system""" - rollback_result = subprocess.Popen( - ["systemctl", "start", "sp-nixos-rollback.service"], start_new_session=True - ) - rollback_result.communicate()[0] - return rollback_result.returncode + run_blocking(["systemctl", "start", "sp-nixos-rollback.service"], new_session=True) + return 0 def upgrade_system() -> int: """Upgrade the system""" - upgrade_result = subprocess.Popen( - ["systemctl", "start", "sp-nixos-upgrade.service"], start_new_session=True - ) - upgrade_result.communicate()[0] - return upgrade_result.returncode + run_blocking(["systemctl", "start", "sp-nixos-upgrade.service"], new_session=True) + return 0 def reboot_system() -> None: """Reboot the system""" - subprocess.Popen(["reboot"], start_new_session=True) + run_blocking(["reboot"], new_session=True) def get_system_version() -> str: diff --git a/selfprivacy_api/graphql/mutations/system_mutations.py b/selfprivacy_api/graphql/mutations/system_mutations.py index b0cdae8..13ac16b 100644 --- a/selfprivacy_api/graphql/mutations/system_mutations.py +++ b/selfprivacy_api/graphql/mutations/system_mutations.py @@ -115,39 +115,67 @@ class SystemMutations: @strawberry.mutation(permission_classes=[IsAuthenticated]) def run_system_rebuild(self) -> GenericMutationReturn: - system_actions.rebuild_system() - return GenericMutationReturn( - success=True, - message="Starting rebuild system", - code=200, - ) + try: + system_actions.rebuild_system() + return GenericMutationReturn( + success=True, + message="Starting rebuild system", + code=200, + ) + except system_actions.ShellException as e: + return GenericMutationReturn( + success=False, + message=str(e), + code=500, + ) @strawberry.mutation(permission_classes=[IsAuthenticated]) def run_system_rollback(self) -> GenericMutationReturn: system_actions.rollback_system() - return GenericMutationReturn( - success=True, - message="Starting rebuild system", - code=200, - ) + try: + return GenericMutationReturn( + success=True, + message="Starting rebuild system", + code=200, + ) + except system_actions.ShellException as e: + return GenericMutationReturn( + success=False, + message=str(e), + code=500, + ) @strawberry.mutation(permission_classes=[IsAuthenticated]) def run_system_upgrade(self) -> GenericMutationReturn: system_actions.upgrade_system() - return GenericMutationReturn( - success=True, - message="Starting rebuild system", - code=200, - ) + try: + return GenericMutationReturn( + success=True, + message="Starting rebuild system", + code=200, + ) + except system_actions.ShellException as e: + return GenericMutationReturn( + success=False, + message=str(e), + code=500, + ) @strawberry.mutation(permission_classes=[IsAuthenticated]) def reboot_system(self) -> GenericMutationReturn: system_actions.reboot_system() - return GenericMutationReturn( - success=True, - message="System reboot has started", - code=200, - ) + try: + return GenericMutationReturn( + success=True, + message="System reboot has started", + code=200, + ) + except system_actions.ShellException as e: + return GenericMutationReturn( + success=False, + message=str(e), + code=500, + ) @strawberry.mutation(permission_classes=[IsAuthenticated]) def pull_repository_changes(self) -> GenericMutationReturn: diff --git a/tests/test_graphql/test_system_nixos_tasks.py b/tests/test_graphql/test_system_nixos_tasks.py index b292fda..6052e9f 100644 --- a/tests/test_graphql/test_system_nixos_tasks.py +++ b/tests/test_graphql/test_system_nixos_tasks.py @@ -23,15 +23,6 @@ class ProcessMock: returncode = 0 -class BrokenServiceMock(ProcessMock): - """Mock subprocess.Popen for broken service""" - - def communicate(): # pylint: disable=no-method-argument - return (b"Testing error", None) - - returncode = 3 - - @pytest.fixture def mock_subprocess_popen(mocker): mock = mocker.patch("subprocess.Popen", autospec=True, return_value=ProcessMock)