feat: Add script-based custom checks in addition to the predefined ch…

…ecks

README.md

@@ -107,6 +107,25 @@ See [`topgrade`](https://github.com/topgrade-rs/topgrade)'s GitHub for configuri
 
 * `network_not_metered`: if true, checks if the current network connection is not marked as metered
 
+### Section: `checks.scripts`
+
+In addition to the predefined checks above, it is also possible to implement
+custom conditions through user-provided scripts and their exit codes.
+Each entry in the `checks.scripts` array must specify the following settings:
+
+* `shell`: specifies the shell used to execute the custom script (e.g. `bash`)
+
+* `run`: specifies the script text to be run using the specified shell
+
+* `message`: an optional message that is shown when the check fails
+
+* `name`: an optional human-readable name for this check
+
+The custom script should use its exit code to indicate whether the updater should proceed
+(`exit code = 0`) or whether updates should be inhibited right now (any non-0 exit code).
+If `message` is not specified but the script has written text to `stdout`,
+that text will be used as the message.
+
 ### Section: `notify`
 
 * `dbus_notify`: enable graphical notifications via dbus
@@ -119,6 +138,35 @@ See [`topgrade`](https://github.com/topgrade-rs/topgrade)'s GitHub for configuri
     max_cpu_load_percent = 50.0 #     CPU Usage <= 50%?
     max_mem_percent = 90.0      #     RAM Usage <= 90%?
     network_not_metered = true  # Abort if network connection is metered
+
+    [[checks.scripts]]
+        name = "Example script that always fails"
+        shell = "bash"
+        run = "exit 1"
+        message = "Failure message - this message will always appear"
+
+    [[checks.scripts]]
+        name = "Example script that always succeeds"
+        shell = "bash"
+        run = "exit 0"
+        message = "Failure message - this message will never appear"
+
+    [[checks.scripts]]
+        name = "Example multiline script with custom message"
+        shell = "bash"
+        run = """
+echo "This is a custom message"
+exit 1
+"""
+
+    [[checks.scripts]]
+        name = "Python script"
+        shell = "python3"
+        run = """
+print("Python also works when installed")
+exit(1)
+"""
+
 [notify]
     dbus_notify = false         # Do not show notifications
 ```

src/ublue_update/cli.py

@@ -9,6 +9,7 @@
 )
 from ublue_update.update_checks.wait import transaction_wait
 from ublue_update.update_inhibitors.hardware import check_hardware_inhibitors
+from ublue_update.update_inhibitors.custom import check_custom_inhibitors
 from ublue_update.config import cfg
 from ublue_update.session import get_xdg_runtime_dir, get_active_sessions
 from ublue_update.filelock import acquire_lock, release_lock
@@ -73,12 +74,12 @@ def ask_for_updates(system):
         run_updates(system, True)
 
 
-def hardware_inhibitor_checks_failed(
+def inhibitor_checks_failed(
     failures: list, hardware_check: bool, system_update_available: bool, system: bool
 ):
     # ask if an update can be performed through dbus notifications
     if system_update_available and not hardware_check:
-        log.info("Harware checks failed, but update is available")
+        log.info("Precondition checks failed, but update is available")
         ask_for_updates(system)
     # notify systemd that the checks have failed,
     # systemd will try to rerun the unit
@@ -230,9 +231,14 @@ def main():
 
     system_update_available: bool = system_update_check()
     if not cli_args.force and not cli_args.updatecheck:
-        hardware_checks_failed, failures = check_hardware_inhibitors()
-        if hardware_checks_failed:
-            hardware_inhibitor_checks_failed(
+        hw_checks_failed, hw_failures = check_hardware_inhibitors()
+        cs_checks_failed, cs_failures = check_custom_inhibitors()
+
+        checks_failed = hw_checks_failed or cs_checks_failed
+        failures = hw_failures + cs_failures
+
+        if checks_failed:
+            inhibitor_checks_failed(
                 failures,
                 cli_args.check,
                 system_update_available,

src/ublue_update/config.py

@@ -38,6 +38,7 @@ class Config:
     min_battery_percent: Optional[float]
     max_cpu_load_percent: Optional[float]
     max_mem_percent: Optional[float]
+    custom_check_scripts: List[dict]
 
     def load_config(self, path=None):
         config_path = path or find_default_config_file()
@@ -51,6 +52,7 @@ def load_values(self, config):
         self.min_battery_percent = load_value(config, "checks", "min_battery_percent")
         self.max_cpu_load_percent = load_value(config, "checks", "max_cpu_load_percent")
         self.max_mem_percent = load_value(config, "checks", "max_mem_percent")
+        self.custom_check_scripts = load_value(config, "checks", "scripts") or []
 
 
 cfg = Config()

src/ublue_update/update_inhibitors/custom.py

@@ -0,0 +1,68 @@
+import psutil
+import subprocess
+from typing import List, Optional
+from logging import getLogger
+from ublue_update.config import cfg
+
+"""Setup logging"""
+log = getLogger(__name__)
+
+
+def run_custom_check_script(script) -> dict:
+    if script.get('shell') != 'bash':
+        raise Exception('checks.scripts.*.shell must be set to \'bash\'')
+
+    log.debug(f"Running script {script}")
+
+    # Run the specified custom script via bash
+    script_result = subprocess.run(
+        ['bash', '-c', script['run']],
+        capture_output=True,
+        check=False,
+    )
+
+    # An exit code of 0 means "OK", a non-zero exit code
+    # means "Do not download or perform updates right now"
+    script_pass: bool = script_result.returncode == 0
+
+    # Use either the message specified in the config,
+    # the output of the script (if not empty), or a fallback
+    script_output: Optional[str] = script_result.stdout.strip()
+    if len(script_output) == 0:
+        script_output = None
+
+    # Write error messages to our log in case of failure
+    # to catch any interpreter errors etc.
+    script_stderr = script_result.stderr.strip()
+    if not script_pass and len(script_stderr) > 0:
+        log.warning(f"A custom check script failed and wrote the following to STDERR:\n====\n{script_stderr}\n====")
+
+    fallback_message = "A custom check script returned a non-0 exit code"
+    script_message = script.get('message') or script_output or fallback_message
+
+    return {
+        "passed": script_pass,
+        "message": script_message,
+    }
+
+
+def run_custom_check_scripts() -> List[dict]:
+    results = []
+    for script in (cfg.custom_check_scripts or []):
+        results.append(run_custom_check_script(script))
+    return results
+
+
+def check_custom_inhibitors() -> bool:
+
+    custom_inhibitors = run_custom_check_scripts()
+
+    failures = []
+    custom_checks_failed = False
+    for inhibitor_result in custom_inhibitors:
+        if not inhibitor_result["passed"]:
+            custom_checks_failed = True
+            failures.append(inhibitor_result["message"])
+    if not custom_checks_failed:
+        log.info("System passed custom checks")
+    return custom_checks_failed, failures