feat(control): add stream-control/agent/netstate
This commit is contained in:
parent
a52a4b1373
commit
9d6b302a16
1 changed files with 153 additions and 0 deletions
153
stream-control/agent/netstate
Normal file
153
stream-control/agent/netstate
Normal file
|
|
@ -0,0 +1,153 @@
|
|||
#!/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)
|
||||
Loading…
Add table
Add a link
Reference in a new issue