feature(system): better error handling for shell calls

remove-rest
Houkime 2023-12-29 13:11:03 +00:00
parent dcf6dd9ac5
commit 4a580e9b7b
3 changed files with 86 additions and 47 deletions

View File

@ -2,7 +2,7 @@
import os import os
import subprocess import subprocess
import pytz import pytz
from typing import Optional from typing import Optional, List
from pydantic import BaseModel from pydantic import BaseModel
from selfprivacy_api.utils import WriteUserData, ReadUserData from selfprivacy_api.utils import WriteUserData, ReadUserData
@ -58,36 +58,56 @@ def set_auto_upgrade_settings(
user_data["autoUpgrade"]["allowReboot"] = allowReboot 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: def rebuild_system() -> int:
"""Rebuild the system""" """Rebuild the system"""
rebuild_result = subprocess.Popen( run_blocking(["systemctl", "start", "sp-nixos-rebuild.service"], new_session=True)
["systemctl", "start", "sp-nixos-rebuild.service"], start_new_session=True return 0
)
rebuild_result.communicate()[0]
return rebuild_result.returncode
def rollback_system() -> int: def rollback_system() -> int:
"""Rollback the system""" """Rollback the system"""
rollback_result = subprocess.Popen( run_blocking(["systemctl", "start", "sp-nixos-rollback.service"], new_session=True)
["systemctl", "start", "sp-nixos-rollback.service"], start_new_session=True return 0
)
rollback_result.communicate()[0]
return rollback_result.returncode
def upgrade_system() -> int: def upgrade_system() -> int:
"""Upgrade the system""" """Upgrade the system"""
upgrade_result = subprocess.Popen( run_blocking(["systemctl", "start", "sp-nixos-upgrade.service"], new_session=True)
["systemctl", "start", "sp-nixos-upgrade.service"], start_new_session=True return 0
)
upgrade_result.communicate()[0]
return upgrade_result.returncode
def reboot_system() -> None: def reboot_system() -> None:
"""Reboot the system""" """Reboot the system"""
subprocess.Popen(["reboot"], start_new_session=True) run_blocking(["reboot"], new_session=True)
def get_system_version() -> str: def get_system_version() -> str:

View File

@ -115,39 +115,67 @@ class SystemMutations:
@strawberry.mutation(permission_classes=[IsAuthenticated]) @strawberry.mutation(permission_classes=[IsAuthenticated])
def run_system_rebuild(self) -> GenericMutationReturn: def run_system_rebuild(self) -> GenericMutationReturn:
system_actions.rebuild_system() try:
return GenericMutationReturn( system_actions.rebuild_system()
success=True, return GenericMutationReturn(
message="Starting rebuild system", success=True,
code=200, 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]) @strawberry.mutation(permission_classes=[IsAuthenticated])
def run_system_rollback(self) -> GenericMutationReturn: def run_system_rollback(self) -> GenericMutationReturn:
system_actions.rollback_system() system_actions.rollback_system()
return GenericMutationReturn( try:
success=True, return GenericMutationReturn(
message="Starting rebuild system", success=True,
code=200, 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]) @strawberry.mutation(permission_classes=[IsAuthenticated])
def run_system_upgrade(self) -> GenericMutationReturn: def run_system_upgrade(self) -> GenericMutationReturn:
system_actions.upgrade_system() system_actions.upgrade_system()
return GenericMutationReturn( try:
success=True, return GenericMutationReturn(
message="Starting rebuild system", success=True,
code=200, 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]) @strawberry.mutation(permission_classes=[IsAuthenticated])
def reboot_system(self) -> GenericMutationReturn: def reboot_system(self) -> GenericMutationReturn:
system_actions.reboot_system() system_actions.reboot_system()
return GenericMutationReturn( try:
success=True, return GenericMutationReturn(
message="System reboot has started", success=True,
code=200, 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]) @strawberry.mutation(permission_classes=[IsAuthenticated])
def pull_repository_changes(self) -> GenericMutationReturn: def pull_repository_changes(self) -> GenericMutationReturn:

View File

@ -23,15 +23,6 @@ class ProcessMock:
returncode = 0 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 @pytest.fixture
def mock_subprocess_popen(mocker): def mock_subprocess_popen(mocker):
mock = mocker.patch("subprocess.Popen", autospec=True, return_value=ProcessMock) mock = mocker.patch("subprocess.Popen", autospec=True, return_value=ProcessMock)