core/tools: denylist power-management subcommands and privileged programs

src/plyngent/tools/__init__.py

@@ -88,6 +88,7 @@ from .workspace import (
 )
 from .workspace import DEFAULT_POLICY_CONFIRM_TIMEOUT_SECONDS as DEFAULT_POLICY_CONFIRM_TIMEOUT_SECONDS
 from .workspace import AccessMode as AccessMode
+from .workspace import DenylistHit as DenylistHit
 from .workspace import WorkspaceError as WorkspaceError
 from .workspace import WorkspacePolicy as WorkspacePolicy
 from .workspace import active_workspace_policy as active_workspace_policy
@@ -104,6 +105,7 @@ from .workspace import get_workspace_root as get_workspace_root
 from .workspace import grant_policy_command as grant_policy_command
 from .workspace import granted_access_mode as granted_access_mode
 from .workspace import list_workspace_allowlist as list_workspace_allowlist
+from .workspace import match_command_denylist as match_command_denylist
 from .workspace import mode_covers as mode_covers
 from .workspace import parse_access_mode as parse_access_mode
 from .workspace import pop_owned_temporary_workspaces as pop_owned_temporary_workspaces

src/plyngent/tools/workspace.py

@@ -5,7 +5,7 @@ from enum import IntEnum
 from pathlib import Path
 from typing import TYPE_CHECKING, cast
 
-from plyngent.tools.command_scan import unwrap_command
+from plyngent.tools.command_scan import basename, unwrap_command
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Mapping, Sequence
@@ -18,6 +18,10 @@ DEFAULT_COMMAND_DENYLIST: frozenset[str] = frozenset(
         "su",
         "doas",
         "pkexec",
+        # Privileged systemd entry points: run0 is sudo's replacement,
+        # systemd-run executes a command as a transient system unit.
+        "run0",
+        "systemd-run",
         "rm",
         "rmdir",
         "mkfs",
@@ -33,6 +37,18 @@ DEFAULT_COMMAND_DENYLIST: frozenset[str] = frozenset(
         "chown",
         "mount",
         "umount",
+        "init",
+        "telinit",
+        "kexec",
+        "pm-suspend",
+        "pm-hibernate",
+        "s2disk",
+        "s2ram",
+        # Machine power state: only these systemctl/loginctl words are blocked,
+        # so read-only calls (``systemctl status``) stay allowed.
+        "systemctl reboot poweroff halt suspend hibernate hybrid-sleep "
+        "suspend-then-hibernate soft-reboot kexec isolate rescue emergency default",
+        "loginctl terminate-session terminate-user kill-session kill-user reboot poweroff soft-reboot",
     }
 )
 
@@ -405,17 +421,88 @@ def argv_shape_error(argv: object) -> str | None:
     return None
 
 
+@dataclass(frozen=True)
+class DenylistHit:
+    """A command denylist entry that a command matches.
+
+    ``program`` names the program to prompt for (and grant by); ``keyword`` is
+    the argument word that made a subcommand entry dangerous, or ``None`` for an
+    entry that blocks the whole program.
+    """
+
+    entry: str
+    program: str
+    keyword: str | None = None
+
+
+def match_command_denylist(argv: Sequence[str], denylist: frozenset[str]) -> list[DenylistHit]:
+    """Return the denylist entries *argv* matches, outermost program first.
+
+    An entry is a program basename, optionally followed by the argument words
+    that make it dangerous (``systemctl reboot poweroff …``). A whole-program
+    entry matches any program of the wrapper chain; a subcommand entry matches
+    only the program that runs, whose arguments are searched for one of its
+    words — so ``systemctl reboot`` and ``systemctl start reboot.target`` are
+    blocked while ``systemctl status`` is not.
+    """
+    scan = unwrap_command(argv)
+    words = _argument_words(scan.command)
+    hits: list[DenylistHit] = []
+    seen: set[str] = set()
+    for name in (*scan.wrappers, scan.base):
+        hit = _denylist_hit(name, base=scan.base, words=words, denylist=denylist)
+        if hit is not None and hit.program not in seen:
+            seen.add(hit.program)
+            hits.append(hit)
+    return hits
+
+
+def _denylist_hit(
+    name: str,
+    *,
+    base: str,
+    words: frozenset[str],
+    denylist: frozenset[str],
+) -> DenylistHit | None:
+    """First denylist entry matching program *name*, if any."""
+    for entry in sorted(denylist):
+        parts = entry.split()
+        if not parts or basename(parts[0]) != name:
+            continue
+        if len(parts) == 1:
+            return DenylistHit(entry=entry, program=name)
+        if name != base:
+            # Argument words describe the program that runs, not a wrapper.
+            continue
+        keyword = next((word for word in parts[1:] if word.lower() in words), None)
+        if keyword is not None:
+            return DenylistHit(entry=entry, program=name, keyword=keyword.lower())
+    return None
+
+
+def _argument_words(command: Sequence[str]) -> frozenset[str]:
+    """Lower-cased argument words of *command*, plus ``*.target`` stems."""
+    words: set[str] = set()
+    for part in command[1:]:
+        token = part.lower()
+        words.add(token)
+        if token.endswith(".target"):
+            words.add(token.removesuffix(".target"))
+    return frozenset(words)
+
+
 def check_command_allowed(argv: list[str]) -> None:
-    """Raise if argv is empty or a basename in the command chain is denylisted.
+    """Raise if argv is empty or any program in the command chain is denylisted.
 
     The chain is the program that really runs plus every wrapper it is reached
     through (``env FOO=1 rm``, ``pdm run rm``, ``sudo -u x rm``), so a denylisted
     program cannot hide behind one. Basenames are compared lower-cased, so
-    ``RM`` and ``rm.exe`` count as ``rm``.
+    ``RM`` and ``rm.exe`` count as ``rm``; entries that list dangerous arguments
+    (``systemctl reboot …``) match those words in the program's own arguments.
 
-    Denylisted basenames are not hard-rejected when a policy confirm hook is
+    Denylisted programs are not hard-rejected when a policy confirm hook is
     installed: the human is asked (with a timeout; default deny). Session grants
-    skip re-prompting for the same basename. Independent of YOLO soft-confirm.
+    skip re-prompting for the same program. Independent of YOLO soft-confirm.
     """
     shape_error = argv_shape_error(argv)
     if shape_error is not None:
@@ -425,13 +512,14 @@ def check_command_allowed(argv: list[str]) -> None:
         msg = "command argv must not be empty"
         raise WorkspaceError(msg)
     policy = active_workspace_policy()
-    scan = unwrap_command(argv)
-    for binary in (*scan.wrappers, scan.base):
-        if binary not in policy.command_denylist or binary in policy.policy_allowed_commands:
+    for hit in match_command_denylist(argv, policy.command_denylist):
+        binary = hit.program
+        if binary in policy.policy_allowed_commands:
             continue
+        why = f": {hit.keyword!r}" if hit.keyword else ""
         hook = policy.policy_confirm_hook
         if hook is None:
-            msg = f"command denied by policy (basename {binary!r} is blocked)"
+            msg = f"command denied by policy (basename {binary!r} is blocked{why})"
             raise WorkspaceError(msg)
         timeout = policy.policy_confirm_timeout_seconds
         try:
@@ -441,7 +529,7 @@ def check_command_allowed(argv: list[str]) -> None:
             raise WorkspaceError(msg) from exc
         if not allowed:
             msg = (
-                f"command denied by policy (basename {binary!r} is blocked; "
+                f"command denied by policy (basename {binary!r} is blocked{why}; "
                 f"user declined or timed out after {timeout:g}s)"
             )
             raise WorkspaceError(msg)

tests/test_tools/test_workspace.py

@@ -277,6 +277,62 @@ def test_command_denylist_prompts_for_the_wrapped_program(workspace: object) ->
         clear_policy_allowed_commands()
 
 
+def test_match_command_denylist_entries() -> None:
+    from plyngent.tools import match_command_denylist
+
+    denylist = frozenset({"rm", "systemctl reboot poweroff", "loginctl kill-session"})
+    hits = match_command_denylist(["sudo", "rm", "-rf", "/"], denylist)
+    assert [(hit.program, hit.keyword) for hit in hits] == [("rm", None)]
+    assert [hit.program for hit in match_command_denylist(["env", "FOO=1", "rm", "x"], denylist)] == ["rm"]
+    hits = match_command_denylist(["sudo", "sudo", "rm", "x"], frozenset({"sudo", "rm"}))
+    assert [hit.program for hit in hits] == ["sudo", "rm"]
+    # Subcommand entries match the program's own arguments, targets included.
+    hits = match_command_denylist(["systemctl", "reboot"], denylist)
+    assert [(hit.program, hit.keyword, hit.entry) for hit in hits] == [
+        ("systemctl", "reboot", "systemctl reboot poweroff")
+    ]
+    assert match_command_denylist(["systemctl", "start", "poweroff.target"], denylist)[0].keyword == "poweroff"
+    assert match_command_denylist(["pdm", "run", "loginctl", "kill-session"], denylist)[0].program == "loginctl"
+    # Read-only and unrelated calls stay clean.
+    assert match_command_denylist(["systemctl", "status", "sshd"], denylist) == []
+    assert match_command_denylist(["loginctl", "list-sessions"], denylist) == []
+
+
+def test_command_denylist_power_subcommands(workspace: object) -> None:
+    del workspace
+    from plyngent.tools.workspace import clear_policy_allowed_commands, set_policy_confirm_hook
+
+    set_policy_confirm_hook(None)
+    clear_policy_allowed_commands()
+    try:
+        with pytest.raises(WorkspaceError, match="basename 'systemctl' is blocked: 'reboot'"):
+            check_command_allowed(["systemctl", "reboot"])
+        with pytest.raises(WorkspaceError, match="basename 'systemctl' is blocked: 'reboot'"):
+            check_command_allowed(["systemctl", "start", "reboot.target"])
+        with pytest.raises(WorkspaceError, match="basename 'systemctl' is blocked: 'poweroff'"):
+            check_command_allowed(["systemctl", "--user", "poweroff"])
+        with pytest.raises(WorkspaceError, match="basename 'systemctl' is blocked: 'emergency'"):
+            check_command_allowed(["timeout", "3", "systemctl", "emergency"])
+        with pytest.raises(WorkspaceError, match="basename 'loginctl' is blocked: 'kill-session'"):
+            check_command_allowed(["loginctl", "kill-session", "3"])
+        with pytest.raises(WorkspaceError, match="basename 'pm-suspend' is blocked"):
+            check_command_allowed(["pm-suspend"])
+        with pytest.raises(WorkspaceError, match="basename 'run0' is blocked"):
+            check_command_allowed(["run0", "id"])
+        with pytest.raises(WorkspaceError, match="basename 'systemd-run' is blocked"):
+            check_command_allowed(["systemd-run", "--user", "echo", "hi"])
+        # The outermost denylisted program is reported first.
+        with pytest.raises(WorkspaceError, match="basename 'sudo' is blocked"):
+            check_command_allowed(["sudo", "systemctl", "reboot"])
+        # Read-only systemd calls stay allowed.
+        check_command_allowed(["systemctl", "status", "sshd"])
+        check_command_allowed(["systemctl", "list-units", "--type=service"])
+        check_command_allowed(["systemctl", "show", "-p", "CanReboot"])
+        check_command_allowed(["loginctl", "list-sessions"])
+    finally:
+        clear_policy_allowed_commands()
+
+
 def test_root_required() -> None:
     from plyngent.tools.context import InstanceState, bind_instance