#!/usr/bin/env python3 """Restricted network-profile helper for the edge host.""" from __future__ import annotations import json, os, re, shlex, sqlite3, subprocess, sys, time DB = "/var/lib/streamscope/control.db" DEV = "ens3" CLASSES = {"ch": "1:200", "vpn": "1:300", "default": "1:999"} SERVICES = ("netqos", "streamscope-worker", "streamscope-admin", "promtail") def run(argv, timeout=30): try: p = subprocess.run(argv, capture_output=True, text=True, timeout=timeout) return p.returncode, p.stdout.strip(), p.stderr.strip() except Exception as exc: return 125, "", str(exc) def db(): c = sqlite3.connect(DB, timeout=15) c.row_factory = sqlite3.Row c.execute("PRAGMA busy_timeout=15000") return c def rate_map(): rc, out, _ = run(["sudo", "tc", "class", "show", "dev", DEV]) found = {} for line in out.splitlines(): m = re.search(r"^class htb (1:\w+).*?\brate\s+([0-9.]+)([KMG])bit", line, re.I) if not m: continue value = float(m.group(2)) * {"K": .001, "M": 1, "G": 1000}[m.group(3).upper()] target = next((k for k, v in CLASSES.items() if v == m.group(1)), m.group(1)) found[target] = value return found, rc def state(): c = db() limits = {r["target"]: r["limit_mbit"] for r in c.execute("SELECT target,limit_mbit FROM global_limits")} settings = {r["key"]: r["value"] for r in c.execute("SELECT key,value FROM settings")} active_sessions = c.execute("SELECT count(*) FROM session_history WHERE active=1").fetchone()[0] recent = [dict(r) for r in c.execute("SELECT ts,user_name,kind,detail FROM multisession_events ORDER BY id DESC LIMIT 8")] active_blocks = c.execute("SELECT count(*) FROM multisession_blocks WHERE active=1 AND expires_at>strftime('%s','now')").fetchone()[0] c.close() tc, tc_rc = rate_map() ipsets = {} for name in ("geo-a-v4", "geo-a-v6", "cloud-v4", "cloud-v6", "persist-v4", "persist-v6"): rc, out, _ = run(["sudo", "ipset", "list", name, "-t"]) m = re.search(r"Number of entries:\s*(\d+)", out) ipsets[name] = int(m.group(1)) if m else 0 services = {} for name in SERVICES: rc, out, _ = run(["systemctl", "is-active", name]) services[name] = out or "unknown" rc, rules, _ = run(["sudo", "iptables", "-t", "mangle", "-S"]) return { "ok": tc_rc == 0 and all(v == "active" for k, v in services.items() if k != "netqos"), "ts": int(time.time()), "limits": limits, "tc": tc, "ipsets": ipsets, "services": services, "active_sessions": active_sessions, "active_blocks": active_blocks, "policy": { "enabled": settings.get("multisession_enabled", "0") == "1", "mode": settings.get("multisession_mode", "monitor"), "min_sessions": int(settings.get("multisession_min_sessions", "2")), "block_minutes": int(settings.get("multisession_block_minutes", "60")), "persist_seconds": int(settings.get("multisession_persist_seconds", "60")), }, "mark_rules": sum(1 for line in rules.splitlines() if "--set-xmark" in line or "--set-mark" in line), "recent_events": recent, } def set_limits(ch, vpn): values = {"ch": float(ch), "vpn": float(vpn)} if any(not 0.1 <= value <= 100 for value in values.values()): raise ValueError("rate outside 0.1..100 Mbit") old, _ = rate_map() changed = [] try: for target, value in values.items(): cls = CLASSES[target] rc, out, err = run(["sudo", "tc", "class", "change", "dev", DEV, "parent", "1:", "classid", cls, "htb", "rate", f"{value:g}mbit", "ceil", f"{value:g}mbit"]) if rc: raise RuntimeError(err or out or f"tc failed for {target}") changed.append(target) c = db() for target, value in values.items(): c.execute("UPDATE global_limits SET limit_mbit=?,updated_at=datetime('now') WHERE target=?", (value, target)) c.execute("UPDATE limit_profiles SET active=0") c.commit(); c.close() except Exception: for target in changed: if target in old: value = old[target]; cls = CLASSES[target] run(["sudo", "tc", "class", "change", "dev", DEV, "parent", "1:", "classid", cls, "htb", "rate", f"{value:g}mbit", "ceil", f"{value:g}mbit"]) raise return {"ok": True, "limits": values, "tc": rate_map()[0]} def set_policy(mode, min_sessions=None, block_minutes=None): if mode not in {"off", "monitor", "block"}: raise ValueError("invalid policy mode") c = db() c.execute("UPDATE settings SET value=? WHERE key='multisession_enabled'", ("0" if mode == "off" else "1",)) if mode != "off": c.execute("UPDATE settings SET value=? WHERE key='multisession_mode'", (mode,)) if min_sessions is not None: value = int(min_sessions) if not 2 <= value <= 8: raise ValueError("min sessions outside 2..8") c.execute("UPDATE settings SET value=? WHERE key='multisession_min_sessions'", (str(value),)) if block_minutes is not None: value = int(block_minutes) if not 1 <= value <= 1440: raise ValueError("block minutes outside 1..1440") c.execute("UPDATE settings SET value=? WHERE key='multisession_block_minutes'", (str(value),)) c.commit(); c.close() rc, out, err = run(["sudo", "systemctl", "restart", "streamscope-worker"], timeout=20) if rc: raise RuntimeError(err or out or "worker restart failed") return {"ok": True, "policy": state()["policy"]} def background_limiter(action): if action not in {"start", "stop", "refresh"}: raise ValueError("invalid limiter action") unit = "netprofile-rebuild" if action in {"start", "refresh"} else "netprofile-clear" rc, out, err = run(["sudo", "systemd-run", "--unit", unit, "--collect", "/usr/lib/systemd/system/netqos/link-optimizer.sh", action], timeout=15) if rc and "already exists" not in (out + err): raise RuntimeError(err or out) return {"ok": True, "started": action} def restart(name): mapping = {"worker": "streamscope-worker", "admin": "streamscope-admin"} if name not in mapping: raise ValueError("invalid service") rc, out, err = run(["sudo", "systemctl", "restart", mapping[name]], timeout=20) if rc: raise RuntimeError(err or out) return {"ok": True, "restarted": name} def main(): raw = os.environ.get("SSH_ORIGINAL_COMMAND", "state").strip() argv = shlex.split(raw) if not argv: argv = ["state"] if argv == ["state"]: result = state() elif len(argv) == 3 and argv[0] == "set-limits": result = set_limits(argv[1], argv[2]) elif argv[0] == "set-policy" and 2 <= len(argv) <= 4: result = set_policy(argv[1], argv[2] if len(argv) > 2 else None, argv[3] if len(argv) > 3 else None) elif len(argv) == 2 and argv[0] == "limiter": result = background_limiter(argv[1]) elif len(argv) == 2 and argv[0] == "restart": result = restart(argv[1]) else: raise ValueError("unsupported command") print(json.dumps(result, ensure_ascii=False, separators=(",", ":"))) if __name__ == "__main__": try: main() except Exception as exc: print(json.dumps({"ok": False, "error": str(exc)[:300]}, ensure_ascii=False, separators=(",", ":"))) sys.exit(1)