From b98b76d4e45392430cf6d468d3cfd6d526ec4657 Mon Sep 17 00:00:00 2001 From: Filip Strajnar Date: Sat, 26 Oct 2024 18:25:37 +0200 Subject: [PATCH] Ensure script is ran as root. --- src/pci_passthrough_assist/__main__.py | 5 +++++ src/pci_passthrough_assist/permissions.py | 5 +++++ src/pci_passthrough_assist/process_runner.py | 14 ++++++++++++++ 3 files changed, 24 insertions(+) create mode 100644 src/pci_passthrough_assist/permissions.py create mode 100644 src/pci_passthrough_assist/process_runner.py diff --git a/src/pci_passthrough_assist/__main__.py b/src/pci_passthrough_assist/__main__.py index b141a85..b57ad6b 100644 --- a/src/pci_passthrough_assist/__main__.py +++ b/src/pci_passthrough_assist/__main__.py @@ -1,5 +1,10 @@ from sys import platform +from pci_passthrough_assist.permissions import is_ran_by_root if platform != "linux": print("This tool will only work on Linux based OS.") + exit(1) + +if not is_ran_by_root(): + print("This script needs to run as root.") exit(1) \ No newline at end of file diff --git a/src/pci_passthrough_assist/permissions.py b/src/pci_passthrough_assist/permissions.py new file mode 100644 index 0000000..7c18e7c --- /dev/null +++ b/src/pci_passthrough_assist/permissions.py @@ -0,0 +1,5 @@ +from pci_passthrough_assist.process_runner import sh + + +def is_ran_by_root() -> bool: + return sh(["whoami"]) == "root" diff --git a/src/pci_passthrough_assist/process_runner.py b/src/pci_passthrough_assist/process_runner.py new file mode 100644 index 0000000..33cb022 --- /dev/null +++ b/src/pci_passthrough_assist/process_runner.py @@ -0,0 +1,14 @@ +from subprocess import run + + +def sh_binary(args: list[str], ) -> bytes: + return run(args, capture_output=True).stdout + + +def sh(args: list[str], rstrip_newline=True): + string_output: str = sh_binary(args).decode() + return string_output if not rstrip_newline else string_output.rstrip("\n") + + +def sh_lines(args: list[str]) -> list[str]: + return sh(args, rstrip_newline=True).splitlines()