Add network_not_metered and custom check scripts (#111)

README.md

@@ -62,6 +62,7 @@ options:
   -c, --check        run update checks and exit
   -u, --updatecheck  check for updates and exit
   -w, --wait         wait for transactions to complete and exit
+  --config CONFIG    use the specified config file
   --system           only run system updates (requires root)
 ```
 
@@ -88,26 +89,95 @@ See [`topgrade`](https://github.com/topgrade-rs/topgrade)'s GitHub for configuri
 
 ## Location
 
-### Valid config paths (in order of priority):
+### Valid config paths (in order of priority from highest to lowest):
 
-```/etc/ublue-update/ublue-update.toml```
+1. ```/etc/ublue-update/ublue-update.toml```
 
-```/usr/etc/ublue-update/ublue-update.toml```
+2. ```/usr/etc/ublue-update/ublue-update.toml```
 
 
 ## Config Variables
-Section: `checks`
+### Section: `checks`
 
-`min_battery_percent`: checks if battery is above specified percent
+* `min_battery_percent`: checks if battery is above specified percent
 
-`max_cpu_load_percent`: checks if cpu average load is under specified percent
+* `max_cpu_load_percent`: checks if cpu average load is under specified percent
 
-`max_mem_percent`: checks if memory usage is below specified the percent
+* `max_mem_percent`: checks if memory usage is below specified percent
 
+* `network_not_metered`: if true, checks if the current network connection is not marked as metered
 
-Section: `notify`
+### Section: `checks.scripts`
 
-`dbus_notify`: enable graphical notifications via dbus
+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 parameters `run` and `file` are mutually exclusive, but at least one must be specified.
+The `shell` parameter is mandatory when using `run`.
+
+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
+
+### Full Example
+
+```toml
+[checks]
+    min_battery_percent = 20.0  # Battery Level >= 20%?
+    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)
+"""
+
+    [[checks.scripts]]
+        name = "Example external script"
+        # shell = "bash" # specifying a shell is optional for external scripts/programs
+        file = "/bin/true"
+
+[notify]
+    dbus_notify = false         # Do not show notifications
+```
 
 ## How do I build this?
 

files/usr/etc/ublue-update/ublue-update.toml

@@ -2,5 +2,6 @@
     min_battery_percent = 50.0
     max_cpu_load_percent = 50.0
     max_mem_percent = 90.0
+    network_not_metered = false
 [notify]
     dbus_notify = true

src/ublue_update/cli.py

@@ -9,13 +9,14 @@
 )
 from ublue_update.update_checks.wait import transaction_wait
 from ublue_update.update_inhibitors.hardware import check_hardware_inhibitors
-from ublue_update.config import load_value
+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
 
 
 def notify(title: str, body: str, actions: list = [], urgency: str = "normal"):
-    if not dbus_notify:
+    if not cfg.dbus_notify:
         return
     process_uid = os.getuid()
     args = [
@@ -58,7 +59,7 @@ def notify(title: str, body: str, actions: list = [], urgency: str = "normal"):
 
 
 def ask_for_updates(system):
-    if not dbus_notify:
+    if not cfg.dbus_notify:
         return
     out = notify(
         "System Updater",
@@ -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
@@ -159,7 +160,7 @@ def run_updates(system, system_update_available):
             )
             log.debug(out.stdout.decode("utf-8"))
         log.info("System update complete")
-        if pending_deployment_check() and system_update_available and dbus_notify:
+        if pending_deployment_check() and system_update_available and cfg.dbus_notify:
             out = notify(
                 "System Updater",
                 "System update complete, pending changes will take effect after reboot. Reboot now?",
@@ -177,12 +178,10 @@ def run_updates(system, system_update_available):
     os._exit(0)
 
 
-dbus_notify: bool = load_value("notify", "dbus_notify")
-
 # setup logging
 logging.basicConfig(
     format="[%(asctime)s] %(name)s:%(levelname)s | %(message)s",
-    level=os.getenv("UBLUE_LOG", default=logging.INFO),
+    level=os.getenv("UBLUE_LOG", default="INFO").upper(),
 )
 log = logging.getLogger(__name__)
 
@@ -212,22 +211,34 @@ def main():
         action="store_true",
         help="wait for transactions to complete and exit",
     )
+    parser.add_argument(
+        "--config",
+        help="use the specified config file"
+    )
     parser.add_argument(
         "--system",
         action="store_true",
         help="only run system updates (requires root)",
     )
     cli_args = parser.parse_args()
-    hardware_checks_failed = False
+
+    # Load the configuration file
+    cfg.load_config(cli_args.config)
 
     if cli_args.wait:
         transaction_wait()
         os._exit(0)
+
     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

@@ -1,34 +1,58 @@
 import tomllib
 import os
+from typing import List, Optional
+from logging import getLogger
 
+"""Setup logging"""
+log = getLogger(__name__)
 
-def load_config():
+
+def find_default_config_file():
     # load config values
     config_paths = [
         "/etc/ublue-update/ublue-update.toml",
         "/usr/etc/ublue-update/ublue-update.toml",
     ]
 
     # search for the right config
-    config_path = ""
-    fallback_config_path = ""
+    # first config file that is found wins
     for path in config_paths:
         if os.path.isfile(path):
-            if config_path == "":
-                config_path = path
-            fallback_config_path = path
-            break
+            return path
+
+    return None
+
+
+def load_value(dct, *keys):
+    for key in keys:
+        try:
+            dct = dct[key]
+        except KeyError:
+            return None
+    return dct
+
 
-    fallback_config = tomllib.load(open(fallback_config_path, "rb"))
-    config = tomllib.load(open(config_path, "rb"))
-    return config, fallback_config
+class Config:
+    dbus_notify: bool
+    network_not_metered: Optional[bool]
+    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()
+        config = tomllib.load(open(config_path, "rb"))
+        log.debug(f"Configuration loaded from {os.path.abspath(config_path)}")
+        self.load_values(config)
 
-def load_value(key, value):
-    fallback = fallback_config[key][value]
-    if key in config.keys():
-        return config[key].get(value, fallback)
-    return fallback
+    def load_values(self, config):
+        self.dbus_notify = load_value(config, "notify", "dbus_notify") or False
+        self.network_not_metered = load_value(config, "checks", "network_not_metered")
+        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 []
 
 
-config, fallback_config = load_config()
+cfg = Config()

src/ublue_update/update_inhibitors/custom.py

@@ -0,0 +1,73 @@
+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 'run' in script and 'shell' not in script:
+        raise Exception('checks.scripts.*: \'shell\' must be specified when \'run\' is used')
+
+    if 'run' in script and 'file' in script:
+        raise Exception('checks.scripts.*: Only one of \'run\' and \'file\' must be set for a given script')
+
+    log.debug(f"Running script {script}")
+
+    # Run the specified custom script
+    if 'run' in script:
+        run_args = [script['shell'], '-c', script['run']]
+    elif 'shell' in script:
+        run_args = [script['shell'], script['file']]
+    else:
+        run_args = [script['file']]
+    script_result = subprocess.run(run_args, capture_output=True, text=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