mitm-webserver/backend/src/api/packet_scripting_api.py

# script_router_stateless.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.
"""

import os
import uuid
import subprocess
import time
import re
import logging
from typing import List, Dict, Optional

from fastapi import APIRouter, UploadFile, File, HTTPException
from fastapi.responses import FileResponse
from pydantic import BaseModel

# ---------- config ----------
SCRIPT_DIR = "/srv/fw-scripts"
UNIT_DIR = "/etc/systemd/system"
UNIT_PREFIX = "fw-script"   # unit names like fw-script-<sid>-q<qnum>.service

os.makedirs(SCRIPT_DIR, exist_ok=True)

# ---------- logging ----------
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s [%(name)s] %(message)s")
logger = logging.getLogger("script-router-stateless")

# ---------- router ----------
router = APIRouter(prefix="/scripts", tags=["scripts"])

# ---------- models ----------
class ScriptInfo(BaseModel):
    id: str
    name: str
    path: str

class EnableRequest(BaseModel):
    qnum: int
    service_name: Optional[str] = None
    extra_args: Optional[str] = None

# ---------- helpers ----------
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."""
    unit_path = unit_path_for(service_name)
    unit_text = f"""[Unit]
Description={description}
After=network.target

[Service]
Type=simple
ExecStart={exec_start}
Restart=always
RestartSec=2
StandardOutput=syslog
StandardError=syslog

[Install]
WantedBy=multi-user.target
"""
    with open(unit_path, "w") as fh:
        fh.write(unit_text)
    # reload systemd to pick up new unit
    subprocess.run(["systemctl", "daemon-reload"], check=True)
    logger.info("Wrote unit %s", unit_path)
    return unit_path

def remove_unit(service_name: str):
    """Stop and remove a unit file, then daemon-reload."""
    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)
    if os.path.exists(unit_path):
        os.remove(unit_path)
        subprocess.run(["systemctl", "daemon-reload"], check=True)
        logger.info("Removed unit %s", unit_path)

def start_unit(service_name: str):
    subprocess.run(["systemctl", "start", service_name], check=True)
    logger.info("Started 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

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
    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]"
_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": None, "script_path": None, "qnum": None, "extra": None}

# ---------- 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)
        raise HTTPException(status_code=400, detail="only .py scripts allowed")
    data = await file.read()
    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}

@router.get("", response_model=List[ScriptInfo])
def list_scripts():
    out = []
    for fn in os.listdir(SCRIPT_DIR):
        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))
    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)
        raise HTTPException(status_code=404, detail="not found")
    logger.info("Download script sid=%s", sid)
    return FileResponse(path, media_type="text/x-python", filename=f"{sid}.py")

@router.post("/{sid}/enable")
def enable_script(sid: str, req: EnableRequest):
    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}"
    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
        time.sleep(0.05)
        start_unit(service_name)
    except subprocess.CalledProcessError as e:
        logger.exception("Failed to start unit %s", service_name)
        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}

@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:
            subprocess.run(["systemctl", "stop", service_name], check=False)
        except Exception:
            pass
        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)
    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 = {}
    for svc in units:
        parsed = parse_unit_execstart(svc)
        active = False
        try:
            active = is_unit_active(svc)
        except Exception:
            active = False
        results[svc] = {"parsed": parsed, "active": active}
    logger.debug("Status queried: found %d fw units", len(results))
    return results

@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:
                active = is_unit_active(svc)
            except Exception:
                active = False
            matches.append({"service": svc, "parsed": parsed, "active": active})
    logger.debug("Status for sid=%s -> %d matches", sid, len(matches))
    return {"sid": sid, "mappings": matches}

# End of router