From f534d025d0fecda8e9311288acdea9c6ab7ca8fa Mon Sep 17 00:00:00 2001 From: Kevin Veen-Birkenbach Date: Fri, 23 Jan 2026 10:23:59 +0100 Subject: [PATCH] feat(persist): add systemd persistence backend with uninstall https://chatgpt.com/share/69733e45-96ec-800f-9aad-0cac7306dedd --- .gitignore | 1 + Makefile | 3 + src/automtu/__init__.py | 2 - src/automtu/cli.py | 12 ++++ src/automtu/core.py | 25 ++++++-- src/automtu/persist.py | 119 ++++++++++++++++++++++++++++++++++++ tests/unit/test___init__.py | 14 ----- tests/unit/test_persist.py | 114 ++++++++++++++++++++++++++++++++++ 8 files changed, 270 insertions(+), 20 deletions(-) create mode 100644 src/automtu/persist.py delete mode 100644 tests/unit/test___init__.py create mode 100644 tests/unit/test_persist.py diff --git a/.gitignore b/.gitignore index 98b37d5..c643e2b 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ *.pyc __pycache__ dist/ +build/* *.egg-info \ No newline at end of file diff --git a/Makefile b/Makefile index 589a23c..877610f 100644 --- a/Makefile +++ b/Makefile @@ -10,3 +10,6 @@ help: test: @PYTHONPATH="$(CURDIR)/src" "$(PYTHON)" -m unittest discover -s tests/unit -p "test_*.py" -v + +install: + sudo pip install -e . --upgrade --break-system-packages diff --git a/src/automtu/__init__.py b/src/automtu/__init__.py index 07c5de9..e69de29 100644 --- a/src/automtu/__init__.py +++ b/src/automtu/__init__.py @@ -1,2 +0,0 @@ -__all__ = ["__version__"] -__version__ = "0.1.0" diff --git a/src/automtu/cli.py b/src/automtu/cli.py index 8efd48d..3565fcd 100644 --- a/src/automtu/cli.py +++ b/src/automtu/cli.py @@ -91,6 +91,18 @@ def build_parser() -> argparse.ArgumentParser: "--dry-run", action="store_true", help="Show actions without applying changes." ) + # --- Persistence --- + ap.add_argument( + "--persist", + choices=["systemd"], + help="Persist MTU configuration across reboots (currently supported: systemd).", + ) + ap.add_argument( + "--uninstall", + action="store_true", + help="Uninstall persistence backend (requires --persist).", + ) + # --- Machine-readable output modes --- ap.add_argument( "--print-mtu", diff --git a/src/automtu/core.py b/src/automtu/core.py index f6fb4bd..ce3e456 100644 --- a/src/automtu/core.py +++ b/src/automtu/core.py @@ -57,14 +57,31 @@ def run_automtu(args) -> int: log = Logger(mode.machine).log - # Root is only needed if we actually change something (without --dry-run). needs_root = bool( - args.apply_egress_mtu - or args.apply_wg_mtu - or (args.force_egress_mtu is not None) + getattr(args, "apply_egress_mtu", False) + or getattr(args, "apply_wg_mtu", False) + or (getattr(args, "force_egress_mtu", None) is not None) + or (getattr(args, "persist", None) is not None) ) require_root(dry=args.dry_run, needs_root=needs_root) + # Persistence mode: install/uninstall persistence mechanism and exit. + if getattr(args, "persist", None): + if args.persist == "systemd": + from .persist import persist_systemd, uninstall_systemd + + if getattr(args, "uninstall", False): + uninstall_systemd(dry=args.dry_run) + return 0 + + persist_systemd(sys.argv, dry=args.dry_run) + return 0 + + print( + f"[automtu][ERROR] Unknown persist backend: {args.persist}", file=sys.stderr + ) + return 4 + egress = args.egress_if or detect_egress_iface(ignore_vpn=not args.prefer_wg_egress) if not egress: print( diff --git a/src/automtu/persist.py b/src/automtu/persist.py new file mode 100644 index 0000000..cf75c21 --- /dev/null +++ b/src/automtu/persist.py @@ -0,0 +1,119 @@ +from __future__ import annotations + +import shlex +import shutil +import subprocess +from pathlib import Path +from typing import List + + +_SYSTEMD_UNIT_PATH = Path("/etc/systemd/system/automtu.service") + + +def _strip_persist_args(argv: List[str]) -> List[str]: + """ + Remove persistence-only arguments from argv: + - --persist systemd + - --persist=systemd + - --uninstall + Keeps all other args as-is. + """ + out: list[str] = [] + i = 0 + while i < len(argv): + a = argv[i] + if a == "--persist": + i += 1 + if i < len(argv) and not argv[i].startswith("-"): + i += 1 + continue + if a.startswith("--persist="): + i += 1 + continue + if a == "--uninstall": + i += 1 + continue + out.append(a) + i += 1 + return out + + +def _resolve_exec(argv0: str) -> str: + """ + Resolve the executable path for systemd ExecStart. + + Prefer an absolute path (from PATH lookup). Fall back to argv0. + """ + resolved = shutil.which(argv0) + if resolved: + return resolved + return argv0 + + +def persist_systemd(argv: List[str], *, dry: bool) -> None: + """ + Install a systemd oneshot service that re-runs automtu with the same arguments. + This provides a distro-agnostic persistence mechanism. + """ + if not argv: + raise ValueError("argv must not be empty") + + filtered = _strip_persist_args(argv[:]) + if not filtered: + raise ValueError("argv filtered to empty; cannot persist") + + exe = _resolve_exec(filtered[0]) + args = [exe, *filtered[1:]] + + execstart = shlex.join(args) + + unit = f"""\ +[Unit] +Description=Auto MTU via automtu +After=network-online.target +Wants=network-online.target + +[Service] +Type=oneshot +ExecStart={execstart} + +[Install] +WantedBy=multi-user.target +""" + + if dry: + print(f"[automtu] DRY-RUN: would write systemd unit to {_SYSTEMD_UNIT_PATH}") + print(unit.rstrip()) + print("[automtu] DRY-RUN: would run: systemctl daemon-reload") + print("[automtu] DRY-RUN: would run: systemctl enable automtu.service") + return + + _SYSTEMD_UNIT_PATH.write_text(unit) + + subprocess.run(["systemctl", "daemon-reload"], check=True) + subprocess.run(["systemctl", "enable", "automtu.service"], check=True) + + print("[automtu] Installed and enabled systemd service: automtu.service") + print("[automtu] Tip: run 'systemctl start automtu.service' to apply immediately.") + + +def uninstall_systemd(*, dry: bool) -> None: + """ + Uninstall the systemd persistence backend: + - disable automtu.service + - remove unit file (if present) + - daemon-reload + """ + if dry: + print("[automtu] DRY-RUN: would run: systemctl disable automtu.service") + print(f"[automtu] DRY-RUN: would remove: {_SYSTEMD_UNIT_PATH} (if exists)") + print("[automtu] DRY-RUN: would run: systemctl daemon-reload") + return + + subprocess.run(["systemctl", "disable", "automtu.service"], check=True) + + if _SYSTEMD_UNIT_PATH.exists(): + _SYSTEMD_UNIT_PATH.unlink() + + subprocess.run(["systemctl", "daemon-reload"], check=True) + print("[automtu] Uninstalled systemd service: automtu.service") diff --git a/tests/unit/test___init__.py b/tests/unit/test___init__.py deleted file mode 100644 index d476ab0..0000000 --- a/tests/unit/test___init__.py +++ /dev/null @@ -1,14 +0,0 @@ -import unittest - -import automtu - - -class TestInit(unittest.TestCase): - def test_version_is_exposed(self) -> None: - self.assertTrue(hasattr(automtu, "__version__")) - self.assertIsInstance(automtu.__version__, str) - self.assertRegex(automtu.__version__, r"^\d+\.\d+\.\d+$") - - -if __name__ == "__main__": - unittest.main(verbosity=2) diff --git a/tests/unit/test_persist.py b/tests/unit/test_persist.py new file mode 100644 index 0000000..ba4cb87 --- /dev/null +++ b/tests/unit/test_persist.py @@ -0,0 +1,114 @@ +import io +import unittest +from contextlib import redirect_stdout +from pathlib import Path +from unittest.mock import patch + +import automtu.persist as persist + + +class TestPersist(unittest.TestCase): + def test_strip_persist_args_removes_flag_and_value(self) -> None: + argv = ["automtu", "--apply-wg-mtu", "--persist", "systemd"] + self.assertEqual( + persist._strip_persist_args(argv), + ["automtu", "--apply-wg-mtu"], + ) + + def test_strip_persist_args_removes_equals_form(self) -> None: + argv = ["automtu", "--apply-wg-mtu", "--persist=systemd"] + self.assertEqual( + persist._strip_persist_args(argv), + ["automtu", "--apply-wg-mtu"], + ) + + def test_strip_persist_args_removes_uninstall(self) -> None: + argv = ["automtu", "--persist", "systemd", "--uninstall", "--apply-wg-mtu"] + self.assertEqual( + persist._strip_persist_args(argv), + ["automtu", "--apply-wg-mtu"], + ) + + def test_persist_systemd_dry_run_prints_unit(self) -> None: + argv = [ + "automtu", + "--auto-pmtu-from-wg", + "--apply-wg-mtu", + "--persist", + "systemd", + ] + + with ( + patch("automtu.persist.shutil.which", return_value="/usr/bin/automtu"), + patch("automtu.persist._SYSTEMD_UNIT_PATH", Path("/tmp/automtu.service")), + ): + out = io.StringIO() + with redirect_stdout(out): + persist.persist_systemd(argv, dry=True) + + s = out.getvalue() + self.assertIn("DRY-RUN", s) + self.assertIn( + "ExecStart=/usr/bin/automtu --auto-pmtu-from-wg --apply-wg-mtu", s + ) + + def test_uninstall_systemd_dry_run_prints_actions(self) -> None: + with patch("automtu.persist._SYSTEMD_UNIT_PATH", Path("/tmp/automtu.service")): + out = io.StringIO() + with redirect_stdout(out): + persist.uninstall_systemd(dry=True) + + s = out.getvalue() + self.assertIn("DRY-RUN", s) + self.assertIn("systemctl disable", s) + + def test_uninstall_systemd_runs_disable_and_removes_unit_when_present(self) -> None: + fake_path = Path("/tmp/automtu.service") + + calls = [] + + def fake_run(cmd, check=False): # type: ignore[no-untyped-def] + calls.append((tuple(cmd), bool(check))) + + with ( + patch("automtu.persist._SYSTEMD_UNIT_PATH", fake_path), + patch("automtu.persist.subprocess.run", side_effect=fake_run), + patch.object(Path, "exists", return_value=True), + patch.object(Path, "unlink", return_value=None), + ): + persist.uninstall_systemd(dry=False) + + # Must disable unit and reload daemon + self.assertIn((("systemctl", "disable", "automtu.service"), True), calls) + self.assertIn((("systemctl", "daemon-reload"), True), calls) + + def test_persist_systemd_writes_unit_and_enables(self) -> None: + argv = ["automtu", "--apply-wg-mtu", "--persist", "systemd"] + + writes = {"text": None} + calls = [] + + def fake_write_text(self, txt, *args, **kwargs): # type: ignore[no-untyped-def] + writes["text"] = txt + return len(txt) + + def fake_run(cmd, check=False): # type: ignore[no-untyped-def] + calls.append((tuple(cmd), bool(check))) + + with ( + patch("automtu.persist.shutil.which", return_value="/usr/bin/automtu"), + patch("automtu.persist._SYSTEMD_UNIT_PATH", Path("/tmp/automtu.service")), + patch.object(Path, "write_text", new=fake_write_text), + patch("automtu.persist.subprocess.run", side_effect=fake_run), + ): + persist.persist_systemd(argv, dry=False) + + self.assertIsNotNone(writes["text"]) + self.assertIn("ExecStart=/usr/bin/automtu --apply-wg-mtu", writes["text"] or "") + + self.assertIn((("systemctl", "daemon-reload"), True), calls) + self.assertIn((("systemctl", "enable", "automtu.service"), True), calls) + + +if __name__ == "__main__": + unittest.main(verbosity=2)