marcus/mitm-webserver
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

scipting improvements
All checks were successful
Build and Deploy MITM Webserver / build (push) Successful in 8s
Details

Browse Source
This commit is contained in:
malmert
2026-01-28 18:25:17 +01:00
parent 0d66ff2694
commit 90087039f7
1 changed files with 226 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files

292
backend/src/api/packet_scripting_api.py
View File

@@ -1,40 +1,56 @@
# script_router_stateless.py
# script_router_with_venv.py
"""
Stateless APIRouter that manages uploaded scripts via systemd only.
- Upload/list/download scripts (stored in SCRIPT_DIR)
- Enable script on qnum => write systemd unit fw-script-<sid>-q<qnum>.service and start it
- Disable script on qnum => stop service and remove unit file
- Status endpoints discover active services by scanning /etc/systemd/system for fw-script-*.service
No DB or run-dir used — systemd is the source of truth.
APIRouter for uploading scripts (with optional requirements.txt),
creating per-script venvs (when requirements provided), and managing
systemd services that run scripts with a queue number argument.
Usage:
from fastapi import FastAPI
from script_router_with_venv import router, register_lifecycle
app = FastAPI()
app.include_router(router)
register_lifecycle(app) # optional: stops/removes manager-created units on shutdown
Notes:
- The router does NOT interact with nft. You should create/delete nft queue rules with your separate API.
- Script files are stored under SCRIPT_DIR.
- Virtualenvs (if created) are stored under VENV_BASE/<sid>.
- Systemd units are created under /etc/systemd/system with names: fw-script-<sid>-q<qnum>.service
- This code must run with permissions to create venvs, write unit files and call systemctl (typically root).
"""
import os
import sys
import uuid
import shutil
import subprocess
import time
import re
import logging
from typing import List, Dict, Optional
from typing import Optional, List, Dict
from fastapi import APIRouter, UploadFile, File, HTTPException
from fastapi.responses import FileResponse
from pydantic import BaseModel
# ---------- config ----------
# ---------- Config ----------
SCRIPT_DIR = "/srv/fw-scripts"
VENV_BASE = "/srv/fw-scripts/venvs"
UNIT_DIR = "/etc/systemd/system"
UNIT_PREFIX = "fw-script" # unit names like fw-script-<sid>-q<qnum>.service
UNIT_PREFIX = "fw-script"
# ensure dirs exist
os.makedirs(SCRIPT_DIR, exist_ok=True)
os.makedirs(VENV_BASE, exist_ok=True)
# ---------- logging ----------
# ---------- Logging ----------
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s [%(name)s] %(message)s")
logger = logging.getLogger("script-router-stateless")
logger = logging.getLogger("script-router-venv")
# ---------- router ----------
# ---------- Router ----------
router = APIRouter(prefix="/scripts", tags=["scripts"])
# ---------- models ----------
# ---------- Models ----------
class ScriptInfo(BaseModel):
id: str
name: str
@@ -44,17 +60,81 @@ class EnableRequest(BaseModel):
qnum: int
service_name: Optional[str] = None
extra_args: Optional[str] = None
enable_at_boot: Optional[bool] = False
# ---------- helpers ----------
# ---------- Utilities: service names / unit paths ----------
def make_service_name(sid: str, qnum: int) -> str:
# keep name safe for systemd
return f"{UNIT_PREFIX}-{sid}-q{qnum}"
def unit_path_for(service_name: str) -> str:
return os.path.join(UNIT_DIR, service_name + ".service")
def write_unit(service_name: str, exec_start: str, description: str = "") -> str:
"""Write unit file and daemon-reload systemd."""
# ---------- Venv helpers ----------
def venv_path_for(sid: str) -> str:
return os.path.join(VENV_BASE, sid)
def venv_python_for(sid: str) -> str:
vpy = os.path.join(venv_path_for(sid), "bin", "python")
if os.path.exists(vpy):
return vpy
# fallback to system python
return "/usr/bin/python3"
def create_venv(sid: str, timeout: int = 60):
venv_dir = venv_path_for(sid)
if os.path.exists(venv_dir):
logger.debug("Venv already exists for sid=%s: %s", sid, venv_dir)
return venv_dir
os.makedirs(os.path.dirname(venv_dir), exist_ok=True)
logger.info("Creating venv for sid=%s at %s", sid, venv_dir)
try:
subprocess.run([sys.executable, "-m", "venv", venv_dir], check=True, timeout=timeout)
except subprocess.CalledProcessError as e:
raise RuntimeError(f"venv creation failed: {e}")
except subprocess.TimeoutExpired:
raise RuntimeError("venv creation timed out")
return venv_dir
def pip_install_requirements(sid: str, requirements_path: str, timeout: int = 600) -> Dict[str, str]:
"""
Install requirements into the venv for sid from requirements_path.
Returns dict: { "stdout": "...", "stderr": "..." }
Raises on failure with details in exception message.
"""
venv_dir = create_venv(sid)
pip_path = os.path.join(venv_dir, "bin", "pip")
# ensure pip exists and upgrade
try:
subprocess.run([pip_path, "install", "--upgrade", "pip"], check=True, capture_output=True, text=True, timeout=300)
except subprocess.CalledProcessError as e:
# continue but warn
logger.warning("pip upgrade failed for sid=%s: %s", sid, e.stderr if hasattr(e, "stderr") else str(e))
# install requirements
try:
p = subprocess.run([pip_path, "install", "-r", requirements_path, "--no-cache-dir"],
check=True, capture_output=True, text=True, timeout=timeout)
logger.info("pip install success for sid=%s", sid)
return {"stdout": p.stdout, "stderr": p.stderr}
except subprocess.CalledProcessError as e:
logger.error("pip install failed for sid=%s: %s", sid, e.stderr if hasattr(e, "stderr") else str(e))
# return output for debugging
out = {"stdout": getattr(e, "stdout", "") or "", "stderr": getattr(e, "stderr", "") or str(e)}
raise RuntimeError(jsonify_cmd_output(out))
except subprocess.TimeoutExpired:
logger.error("pip install timeout for sid=%s", sid)
raise RuntimeError("pip install timed out")
def jsonify_cmd_output(out: Dict[str, str]) -> str:
# helper to pack stdout/stderr into a single string message
s = ""
if out.get("stdout"):
s += "STDOUT:\n" + out["stdout"] + "\n"
if out.get("stderr"):
s += "STDERR:\n" + out["stderr"] + "\n"
return s.strip()
# ---------- systemd helpers ----------
def write_unit(service_name: str, exec_start: str, description: str = "", enable_at_boot: bool = False) -> str:
unit_path = unit_path_for(service_name)
unit_text = f"""[Unit]
Description={description}
@@ -73,19 +153,28 @@ WantedBy=multi-user.target
"""
with open(unit_path, "w") as fh:
fh.write(unit_text)
# reload systemd to pick up new unit
# reload systemd
subprocess.run(["systemctl", "daemon-reload"], check=True)
logger.info("Wrote unit %s", unit_path)
if enable_at_boot:
try:
subprocess.run(["systemctl", "enable", service_name], check=True)
logger.info("Enabled %s at boot", service_name)
except subprocess.CalledProcessError:
logger.warning("Failed to enable %s at boot", service_name)
return unit_path
def remove_unit(service_name: str):
"""Stop and remove a unit file, then daemon-reload."""
unit_path = unit_path_for(service_name)
try:
subprocess.run(["systemctl", "stop", service_name], check=False)
except Exception:
logger.exception("Failed to stop %s", service_name)
unit_path = unit_path_for(service_name)
logger.exception("systemctl stop failed for %s", service_name)
if os.path.exists(unit_path):
try:
subprocess.run(["systemctl", "disable", service_name], check=False)
except Exception:
pass
os.remove(unit_path)
subprocess.run(["systemctl", "daemon-reload"], check=True)
logger.info("Removed unit %s", unit_path)
@@ -94,6 +183,10 @@ def start_unit(service_name: str):
subprocess.run(["systemctl", "start", service_name], check=True)
logger.info("Started service %s", service_name)
def stop_unit(service_name: str):
subprocess.run(["systemctl", "stop", service_name], check=True)
logger.info("Stopped service %s", service_name)
def is_unit_active(service_name: str) -> bool:
p = subprocess.run(["systemctl", "is-active", "--quiet", service_name])
return p.returncode == 0
@@ -101,60 +194,86 @@ def is_unit_active(service_name: str) -> bool:
def list_fw_units() -> List[str]:
"""Return list of fw unit names without .service suffix."""
units = []
# scan UNIT_DIR for files matching UNIT_PREFIX-*.service
try:
for fn in os.listdir(UNIT_DIR):
if fn.startswith(UNIT_PREFIX + "-") and fn.endswith(".service"):
units.append(fn[:-8]) # strip .service
units.append(fn[:-8]) # remove .service
except FileNotFoundError:
logger.warning("Unit dir %s not found", UNIT_DIR)
return units
# ExecStart parser: we expect ExecStart like "/usr/bin/python3 /srv/fw-scripts/<sid>.py <qnum> [extra]"
# ExecStart parse pattern: python path + script path + qnum + optional extra
_RE_EXECSTART = re.compile(r'(?P<py>/\S*python\S*)\s+(?P<script>/\S*?/srv/fw-scripts/(?P<sid>[0-9a-fA-F]+)\.py)\s+(?P<qnum>\d+)(?:\s+(?P<extra>.*))?')
def parse_unit_execstart(service_name: str) -> Optional[Dict]:
"""Read unit file and extract sid and qnum from ExecStart if it matches expected pattern."""
unit_path = unit_path_for(service_name)
if not os.path.exists(unit_path):
return None
with open(unit_path, "r") as fh:
content = fh.read()
# find ExecStart=
m = re.search(r'^ExecStart=(.+)$', content, flags=re.MULTILINE)
if not m:
return None
exec_start = m.group(1).strip()
# try parse
match = _RE_EXECSTART.search(exec_start)
if match:
sd = match.groupdict()
return {
"service": service_name,
"exec_start": exec_start,
"sid": sd["sid"],
"script_path": sd["script"],
"qnum": int(sd["qnum"]),
"extra": sd.get("extra") or ""
}
# fallback: return exec_start only
return {"service": service_name, "exec_start": exec_start, "sid": sd["sid"], "script_path": sd["script"], "qnum": int(sd["qnum"]), "extra": sd.get("extra") or ""}
return {"service": service_name, "exec_start": exec_start, "sid": None, "script_path": None, "qnum": None, "extra": None}
# ---------- endpoints ----------
# ---------- End utilities ----------
# ---------- Endpoints ----------
@router.post("", response_model=ScriptInfo)
async def upload_script(file: UploadFile = File(...)):
if not file.filename.endswith(".py"):
logger.warning("Reject upload (not .py): %s", file.filename)
async def upload_script(script: UploadFile = File(...), requirements: Optional[UploadFile] = File(None)):
"""
Upload a script and optional requirements.txt.
If requirements is supplied, a venv for the script will be created and pip will install the requirements.
Returns sid, path, and pip output if any.
"""
if not script.filename.endswith(".py"):
logger.warning("Rejected upload with invalid extension: %s", script.filename)
raise HTTPException(status_code=400, detail="only .py scripts allowed")
data = await file.read()
data = await script.read()
if len(data) > 2_000_000:
logger.warning("Rejected upload too large: %s size=%d", script.filename, len(data))
raise HTTPException(status_code=400, detail="script too large")
sid = uuid.uuid4().hex
fname = f"{sid}.py"
path = os.path.join(SCRIPT_DIR, fname)
with open(path, "wb") as fh:
fh.write(data)
os.chmod(path, 0o700)
logger.info("Uploaded script %s as %s", file.filename, path)
return {"id": sid, "name": file.filename, "path": path}
logger.info("Uploaded script %s as sid=%s path=%s", script.filename, sid, path)
pip_output = None
if requirements is not None:
# save requirements to a temp path
req_data = await requirements.read()
req_path = os.path.join(SCRIPT_DIR, f"{sid}-requirements.txt")
with open(req_path, "wb") as fh:
fh.write(req_data)
logger.info("Saved requirements for sid=%s at %s (size=%d)", sid, req_path, len(req_data))
# create venv and install
try:
create_venv(sid)
res = pip_install_requirements(sid, req_path)
pip_output = {"stdout": res.get("stdout", ""), "stderr": res.get("stderr", "")}
logger.info("Installed requirements for sid=%s", sid)
except Exception as e:
# cleanup venv on failure (optional)
logger.exception("pip install failed for sid=%s: %s", sid, e)
# Provide useful error to client but keep script stored for inspection
raise HTTPException(status_code=500, detail=f"pip install failed: {e}")
# return pip_output in headers/body? We'll include in body (if present).
resp = {"id": sid, "name": script.filename, "path": path}
if pip_output is not None:
resp["pip"] = pip_output
return resp
@router.get("", response_model=List[ScriptInfo])
def list_scripts():
@@ -163,72 +282,91 @@ def list_scripts():
if not fn.endswith(".py"):
continue
sid = fn.rsplit(".", 1)[0]
out.append({"id": sid, "name": fn, "path": os.path.join(SCRIPT_DIR, fn)})
logger.debug("Listed %d scripts", len(out))
p = os.path.join(SCRIPT_DIR, fn)
out.append({"id": sid, "name": fn, "path": p})
logger.debug("Listed scripts: %d", len(out))
return out
@router.get("/{sid}")
def download_script(sid: str):
path = os.path.join(SCRIPT_DIR, f"{sid}.py")
if not os.path.exists(path):
logger.warning("Download missing script sid=%s", sid)
logger.warning("Download requested for missing sid=%s", sid)
raise HTTPException(status_code=404, detail="not found")
logger.info("Download script sid=%s", sid)
logger.info("Download script sid=%s path=%s", sid, path)
return FileResponse(path, media_type="text/x-python", filename=f"{sid}.py")
@router.post("/{sid}/enable")
def enable_script(sid: str, req: EnableRequest):
"""
Create and start systemd service that runs the script with the queue number.
If a venv was created at upload time, ExecStart will use that venv's python.
"""
script_path = os.path.join(SCRIPT_DIR, f"{sid}.py")
if not os.path.exists(script_path):
logger.warning("Enable requested for missing sid=%s", sid)
raise HTTPException(status_code=404, detail="script not found")
qnum = req.qnum
service_name = req.service_name or make_service_name(sid, qnum)
exec_start = f"/usr/bin/python3 {script_path} {qnum}"
# prefer venv python if exists
python_path = venv_python_for(sid)
exec_start = f"{python_path} {script_path} {qnum}"
if req.extra_args:
exec_start += " " + req.extra_args
# write unit and start
try:
write_unit(service_name, exec_start, description=f"FW script {sid} queue {qnum}")
# small pause
write_unit(service_name, exec_start, description=f"FW script {sid} queue {qnum}", enable_at_boot=req.enable_at_boot)
# slight delay then start
time.sleep(0.05)
start_unit(service_name)
except subprocess.CalledProcessError as e:
logger.exception("Failed to start unit %s", service_name)
logger.exception("Failed to start service %s for sid=%s: %s", service_name, sid, e)
try:
remove_unit(service_name)
except Exception:
pass
raise HTTPException(status_code=500, detail=f"systemd start failed: {e}")
logger.info("Enabled script %s on qnum=%s as service=%s", sid, qnum, service_name)
return {"status": "ok", "sid": sid, "qnum": qnum, "service": service_name}
except Exception as e:
logger.exception("Unknown error starting service %s for sid=%s: %s", service_name, sid, e)
try:
remove_unit(service_name)
except Exception:
pass
raise HTTPException(status_code=500, detail=str(e))
logger.info("Enabled script sid=%s on qnum=%d as service=%s (python=%s)", sid, qnum, service_name, python_path)
return {"status": "ok", "sid": sid, "qnum": qnum, "service": service_name, "python": python_path}
@router.post("/{sid}/disable")
def disable_script(sid: str, qnum: int):
service_name = make_service_name(sid, qnum)
# if unit doesn't exist, still try to remove possible legacy name provided by user
unit_p = unit_path_for(service_name)
if not os.path.exists(unit_p):
logger.warning("Disable requested but unit missing for sid=%s qnum=%s", sid, qnum)
# still attempt to stop in case it exists in systemd (but unit file missing)
# try stopping anyway in case only registered with systemd without unit file
try:
subprocess.run(["systemctl", "stop", service_name], check=False)
except Exception:
pass
logger.warning("Disable requested but unit missing for sid=%s qnum=%s", sid, qnum)
raise HTTPException(status_code=404, detail="service/unit not found")
# stop+remove unit
remove_unit(service_name)
logger.info("Disabled script %s on qnum=%s (removed service %s)", sid, qnum, service_name)
try:
stop_unit(service_name)
except subprocess.CalledProcessError:
logger.exception("Failed stopping service %s", service_name)
try:
remove_unit(service_name)
except Exception:
logger.exception("Failed removing unit %s", service_name)
logger.info("Disabled script sid=%s on qnum=%s (removed service %s)", sid, qnum, service_name)
return {"status": "ok", "sid": sid, "qnum": qnum}
@router.get("/status")
def status_all():
"""Discover all fw-script- units and report status and parsed metadata."""
units = list_fw_units()
results = {}
results: Dict[str, Dict] = {}
for svc in units:
parsed = parse_unit_execstart(svc)
active = False
try:
active = is_unit_active(svc)
except Exception:
@@ -239,11 +377,9 @@ def status_all():
@router.get("/{sid}/status")
def status_for_sid(sid: str):
"""Return status for all units that belong to sid (matching naming convention)."""
units = list_fw_units()
matches = []
for svc in units:
# svc name like fw-script-<sid>-q<qnum>
if svc.startswith(f"{UNIT_PREFIX}-{sid}-q"):
parsed = parse_unit_execstart(svc)
try:
@@ -254,4 +390,28 @@ def status_for_sid(sid: str):
logger.debug("Status for sid=%s -> %d matches", sid, len(matches))
return {"sid": sid, "mappings": matches}
# ---------- Lifecycle helper (optional) ----------
def register_lifecycle(app):
"""
Optionally call this in your main FastAPI app to stop & remove any manager-created units at shutdown.
This will scan for units with our prefix and remove them -- WARNING: if you want systemd to keep services after manager stops,
do NOT register this lifecycle handler.
"""
@app.on_event("shutdown")
def _shutdown_event():
logger.info("Shutdown: removing manager-created systemd units with prefix %s", UNIT_PREFIX)
units = list_fw_units()
for svc in units:
# only remove units that match our naming convention (fw-script-<sid>-q<qnum>)
if svc.startswith(UNIT_PREFIX + "-"):
try:
subprocess.run(["systemctl", "stop", svc], check=False)
except Exception:
pass
try:
remove_unit(svc)
except Exception:
logger.exception("Failed to remove unit %s on shutdown", svc)
logger.info("Shutdown cleanup complete")
# End of router