test scripting changes

backend/src/api/packet_scripting_api.py

@@ -18,11 +18,11 @@ import shutil
 import subprocess
 import time
 import logging
-from typing import Optional, List, Dict
+from typing import Optional, List, Dict, Any
 
 from fastapi import APIRouter, UploadFile, File, Form, HTTPException, Query
 from fastapi.responses import FileResponse
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 # ---------- Configuration ----------
 SCRIPT_DIR = "/srv/fw-scripts"
@@ -44,6 +44,7 @@ router = APIRouter(prefix="/scripts", tags=["scripts"])
 # ---------- Name validation ----------
 _NAME_RE = re.compile(r'^[A-Za-z0-9_.-]+$')
 
+
 def validate_name(name: str) -> None:
     if not name:
         raise ValueError("name must be provided")
@@ -52,35 +53,43 @@ def validate_name(name: str) -> None:
     if name in (".", ".."):
         raise ValueError("invalid name")
 
+
 # ---------- Paths ----------
 def script_path_for(name: str) -> str:
     return os.path.join(SCRIPT_DIR, f"{name}.py")
 
+
 def requirements_path_for(name: str) -> str:
     return os.path.join(SCRIPT_DIR, f"{name}-requirements.txt")
 
+
 def venv_path_for(name: str) -> str:
     return os.path.join(VENV_BASE, name)
 
+
 def venv_python_for(name: str) -> str:
     vpy = os.path.join(venv_path_for(name), "bin", "python")
     if os.path.exists(vpy):
         return vpy
     return "/usr/bin/python3"
 
+
 def make_service_name(name: str, qnum: int) -> str:
     # safe, deterministic service name
     return f"{UNIT_PREFIX}-{name}-q{qnum}"
 
+
 def unit_path_for_name(service_name: str) -> str:
     # default location for units we write
     return os.path.join(UNIT_DIR, service_name + ".service")
 
+
 # ---------- systemd interaction (systemctl-based, no fallback) ----------
 def _systemctl_unit_name(unit: str) -> str:
     """Return unit with .service suffix if missing."""
     return unit if unit.endswith(".service") else unit + ".service"
 
+
 def list_fw_units() -> List[str]:
     """
     Return list of systemd units (without .service suffix) whose name starts with UNIT_PREFIX-.
@@ -88,8 +97,13 @@ def list_fw_units() -> List[str]:
     """
     units: List[str] = []
     try:
-        p = subprocess.run(["systemctl", "list-units", "--type=service", "--all", "--no-legend"],
-                           capture_output=True, text=True, check=False, timeout=3)
+        p = subprocess.run(
+            ["systemctl", "list-units", "--type=service", "--all", "--no-legend"],
+            capture_output=True,
+            text=True,
+            check=False,
+            timeout=3,
+        )
         out = p.stdout or ""
         # each line starts with unit name
         for line in out.splitlines():
@@ -104,6 +118,7 @@ def list_fw_units() -> List[str]:
         logger.exception("systemctl list-units failed")
     return units
 
+
 def get_unit_fragment_path(service_name: str) -> Optional[str]:
     """
     Use `systemctl show -p FragmentPath --value` to obtain the unit file path (if any).
@@ -121,13 +136,28 @@ def get_unit_fragment_path(service_name: str) -> Optional[str]:
         logger.exception("systemctl show FragmentPath failed for %s", service_name)
         return None
 
-def parse_unit_execstart(service_name: str) -> Dict:
+
+class UnitParsed(BaseModel):
+    service: str
+    exec_start: Optional[str] = None
+    name: Optional[str] = None
+    script_path: Optional[str] = None
+    qnum: Optional[int] = None
+    extra: Optional[str] = None
+
+
+class UnitStatus(BaseModel):
+    parsed: UnitParsed
+    active: bool = False
+
+
+def parse_unit_execstart(service_name: str) -> UnitParsed:
     """
     Parser for ExecStart using systemctl --no-pager show.
     - reconstructs argv[] entries if systemd returns a structured blob
     - finds /srv/fw-scripts/<name>.py anywhere in the commandline
     - extracts the first integer after the script as qnum
-    Returns a dict: {service, exec_start, name, script_path, qnum, extra}
+    Returns a UnitParsed instance.
     """
     unit = service_name if service_name.endswith(".service") else service_name + ".service"
 
@@ -142,31 +172,17 @@ def parse_unit_execstart(service_name: str) -> Dict:
         raw = (p.stdout or "").strip()
     except Exception:
         logger.exception("systemctl show ExecStart failed for %s", service_name)
-        return {
-            "service": service_name,
-            "exec_start": None,
-            "name": None,
-            "script_path": None,
-            "qnum": None,
-            "extra": None,
-        }
+        return UnitParsed(service=service_name)
 
     if not raw:
-        return {
-            "service": service_name,
-            "exec_start": None,
-            "name": None,
-            "script_path": None,
-            "qnum": None,
-            "extra": None,
-        }
+        return UnitParsed(service=service_name)
 
-    exec_start = raw
+    exec_start: str = raw
 
     # If structured with argv[] entries, extract them (allow newlines).
     if "argv[]=" in raw:
         argv_entries = re.findall(r'argv\[\]=([^;]+)', raw, flags=re.DOTALL)
-        argv_parts = []
+        argv_parts: List[str] = []
         for a in argv_entries:
             v = a.strip()
             # strip surrounding quotes if present
@@ -181,22 +197,15 @@ def parse_unit_execstart(service_name: str) -> Dict:
     m = re.search(r'(/srv/fw-scripts/(?P<name>[A-Za-z0-9_.-]+)\.py)\b', exec_start)
     if not m:
         # return raw exec_start but no parsed metadata
-        return {
-            "service": service_name,
-            "exec_start": exec_start,
-            "name": None,
-            "script_path": None,
-            "qnum": None,
-            "extra": None,
-        }
+        return UnitParsed(service=service_name, exec_start=exec_start)
 
     script_path = m.group(1)
     name = m.group("name")
 
     # Find first integer token after the script path (queue number)
     after = exec_start[m.end():].strip()
-    qnum = None
-    extra = None
+    qnum: Optional[int] = None
+    extra: Optional[str] = None
     if after:
         tokens = after.split()
         for i, t in enumerate(tokens):
@@ -204,19 +213,12 @@ def parse_unit_execstart(service_name: str) -> Dict:
                 val = int(t)
                 if 0 <= val <= 65535:
                     qnum = val
-                    extra = " ".join(tokens[i+1:]) if i+1 < len(tokens) else ""
+                    extra = " ".join(tokens[i+1:]) if i + 1 < len(tokens) else ""
                     break
             except ValueError:
                 continue
 
-    return {
-        "service": service_name,
-        "exec_start": exec_start,
-        "name": name,
-        "script_path": script_path,
-        "qnum": qnum,
-        "extra": extra,
-    }
+    return UnitParsed(service=service_name, exec_start=exec_start, name=name, script_path=script_path, qnum=qnum, extra=extra)
 
 
 def is_unit_active(service_name: str) -> bool:
@@ -224,9 +226,11 @@ def is_unit_active(service_name: str) -> bool:
     p = subprocess.run(["systemctl", "is-active", "--quiet", unit])
     return p.returncode == 0
 
+
 def write_unit(service_name: str, exec_start: str, description: str = "", enable_at_boot: bool = False) -> str:
     """
     Write unit file to default UNIT_DIR and daemon-reload. This writes to disk as systemd expects.
+    Returns the path to the unit file written.
     """
     unit_path = unit_path_for_name(service_name)
     unit_text = f"""[Unit]
@@ -256,18 +260,22 @@ WantedBy=multi-user.target
             logger.warning("Failed to enable %s at boot", service_name)
     return unit_path
 
+
 def start_unit(service_name: str) -> None:
     subprocess.run(["systemctl", "start", _systemctl_unit_name(service_name)], check=True)
     logger.info("Started service %s", service_name)
 
+
 def stop_unit(service_name: str) -> None:
     subprocess.run(["systemctl", "stop", _systemctl_unit_name(service_name)], check=True)
     logger.info("Stopped service %s", service_name)
 
+
 def disable_unit(service_name: str) -> None:
     subprocess.run(["systemctl", "disable", _systemctl_unit_name(service_name)], check=False)
     logger.info("Disabled service %s", service_name)
 
+
 def remove_unit(service_name: str) -> None:
     """
     Stop + disable the unit, remove the unit file using FragmentPath (if present),
@@ -300,6 +308,7 @@ def remove_unit(service_name: str) -> None:
     except Exception:
         logger.exception("daemon-reload failed after removing unit %s", service_name)
 
+
 # ---------- venv + pip helpers ----------
 def create_venv(name: str, timeout: int = 60) -> str:
     venv_dir = venv_path_for(name)
@@ -318,6 +327,16 @@ def create_venv(name: str, timeout: int = 60) -> str:
         raise RuntimeError("venv creation timed out")
     return venv_dir
 
+
+def jsonify_cmd_output(out: Dict[str, str]) -> str:
+    s = ""
+    if out.get("stdout"):
+        s += "STDOUT:\n" + out["stdout"] + "\n"
+    if out.get("stderr"):
+        s += "STDERR:\n" + out["stderr"] + "\n"
+    return s.strip()
+
+
 def pip_install_requirements(name: str, requirements_path: str, timeout: int = 600) -> Dict[str, str]:
     venv_dir = create_venv(name)
     pip_path = os.path.join(venv_dir, "bin", "pip")
@@ -338,53 +357,72 @@ def pip_install_requirements(name: str, requirements_path: str, timeout: int = 6
         logger.error("pip install timed out for %s", name)
         raise RuntimeError("pip install timed out")
 
-def jsonify_cmd_output(out: Dict[str, str]) -> str:
-    s = ""
-    if out.get("stdout"):
-        s += "STDOUT:\n" + out["stdout"] + "\n"
-    if out.get("stderr"):
-        s += "STDERR:\n" + out["stderr"] + "\n"
-    return s.strip()
 
 # ---------- Models ----------
 class ScriptInfo(BaseModel):
     name: str
     path: str
 
+
+class ScriptUploadResponse(ScriptInfo):
+    pip: Optional[Dict[str, str]] = None
+
+
 class EnableRequest(BaseModel):
-    qnum: int
+    qnum: int = Field(..., description="Queue number (integer)")
     service_name: Optional[str] = None
     extra_args: Optional[str] = None
     enable_at_boot: Optional[bool] = False
 
+
+class UnitMapping(BaseModel):
+    service: str
+    parsed: UnitParsed
+    active: bool = False
+
+
+class StatusForNameResponse(BaseModel):
+    name: str
+    mappings: List[UnitMapping]
+
+
+class OperationResult(BaseModel):
+    status: str
+    name: str
+    qnum: Optional[int] = None
+    service: Optional[str] = None
+    python: Optional[str] = None
+
+
 # ---------- Endpoints ----------
 # Note: Place status endpoints before the dynamic GET /{name} route to avoid routing conflicts.
 
-@router.get("/status")
-def status_all():
+@router.get("/status", response_model=Dict[str, UnitStatus])
+def status_all() -> Dict[str, UnitStatus]:
     """
     Discover all fw-script units via systemctl and report parsed ExecStart + active state.
     """
     units = list_fw_units()
-    results: Dict[str, Dict] = {}
+    results: Dict[str, UnitStatus] = {}
     for svc in units:
         parsed = parse_unit_execstart(svc)
         try:
             active = is_unit_active(svc)
         except Exception:
             active = False
-        results[svc] = {"parsed": parsed, "active": active}
+        results[svc] = UnitStatus(parsed=parsed, active=active)
     logger.debug("Status queried: found %d units", len(results))
     return results
 
-@router.get("/{name}/status")
-def status_for_name(name: str):
+
+@router.get("/{name}/status", response_model=StatusForNameResponse)
+def status_for_name(name: str) -> StatusForNameResponse:
     try:
         validate_name(name)
     except ValueError as e:
         raise HTTPException(status_code=400, detail=str(e))
     units = list_fw_units()
-    matches = []
+    matches: List[UnitMapping] = []
     prefix = f"{UNIT_PREFIX}-{name}-q"
     for svc in units:
         if svc.startswith(prefix):
@@ -393,16 +431,17 @@ def status_for_name(name: str):
                 active = is_unit_active(svc)
             except Exception:
                 active = False
-            matches.append({"service": svc, "parsed": parsed, "active": active})
+            matches.append(UnitMapping(service=svc, parsed=parsed, active=active))
     logger.debug("Status for %s -> %d matches", name, len(matches))
-    return {"name": name, "mappings": matches}
+    return StatusForNameResponse(name=name, mappings=matches)
 
-@router.post("", response_model=ScriptInfo)
+
+@router.post("", response_model=ScriptUploadResponse)
 async def upload_script(
     script: UploadFile = File(...),
     name: str = Form(...),
     requirements: Optional[UploadFile] = File(None),
-):
+) -> ScriptUploadResponse:
     """
     Upload a script with supplied 'name' and optional requirements file.
     On pip/venv install failure, cleanup uploaded files and venv and return 500 with details.
@@ -433,8 +472,8 @@ async def upload_script(
         logger.exception("Failed to write script file for %s: %s", name, e)
         raise HTTPException(status_code=500, detail="failed to save script")
 
-    req_path = None
-    pip_output = None
+    req_path: Optional[str] = None
+    pip_output: Optional[Dict[str, str]] = None
     venv_created = False
 
     try:
@@ -488,19 +527,21 @@ async def upload_script(
     resp = {"name": name, "path": spath}
     if pip_output is not None:
         resp["pip"] = pip_output
-    return resp
+    return ScriptUploadResponse(**resp)
+
 
 @router.get("", response_model=List[ScriptInfo])
-def list_scripts():
-    out = []
+def list_scripts() -> List[ScriptInfo]:
+    out: List[ScriptInfo] = []
     for fn in os.listdir(SCRIPT_DIR):
         if not fn.endswith(".py"):
             continue
         name = fn.rsplit(".", 1)[0]
-        out.append({"name": name, "path": os.path.join(SCRIPT_DIR, fn)})
+        out.append(ScriptInfo(name=name, path=os.path.join(SCRIPT_DIR, fn)))
     logger.debug("Listed %d scripts", len(out))
     return out
 
+
 @router.get("/{name}")
 def download_script(name: str):
     # dynamic route - placed after /status and /{name}/status
@@ -515,8 +556,9 @@ def download_script(name: str):
     logger.info("Download script %s", name)
     return FileResponse(path, media_type="text/x-python", filename=f"{name}.py")
 
-@router.post("/{name}/enable")
-def enable_script(name: str, req: EnableRequest):
+
+@router.post("/{name}/enable", response_model=OperationResult)
+def enable_script(name: str, req: EnableRequest) -> OperationResult:
     try:
         validate_name(name)
     except ValueError as e:
@@ -534,7 +576,7 @@ def enable_script(name: str, req: EnableRequest):
         exec_start += " " + req.extra_args
 
     try:
-        write_unit(service_name, exec_start, description=f"FW script {name} queue {qnum}", enable_at_boot=req.enable_at_boot)
+        write_unit(service_name, exec_start, description=f"FW script {name} queue {qnum}", enable_at_boot=req.enable_at_boot or False)
         time.sleep(0.05)
         start_unit(service_name)
     except subprocess.CalledProcessError as e:
@@ -553,10 +595,11 @@ def enable_script(name: str, req: EnableRequest):
         raise HTTPException(status_code=500, detail=str(e))
 
     logger.info("Enabled script %s on qnum=%d as service=%s (python=%s)", name, qnum, service_name, python_path)
-    return {"status": "ok", "name": name, "qnum": qnum, "service": service_name, "python": python_path}
+    return OperationResult(status="ok", name=name, qnum=qnum, service=service_name, python=python_path)
 
-@router.post("/{name}/disable")
-def disable_script(name: str, qnum: int):
+
+@router.post("/{name}/disable", response_model=OperationResult)
+def disable_script(name: str, qnum: int) -> OperationResult:
     try:
         validate_name(name)
     except ValueError as e:
@@ -576,10 +619,11 @@ def disable_script(name: str, qnum: int):
     except Exception as e:
         logger.exception("Failed to disable unit %s: %s", service_name, e)
         raise HTTPException(status_code=500, detail=str(e))
-    return {"status": "ok", "name": name, "qnum": qnum}
+    return OperationResult(status="ok", name=name, qnum=qnum, service=service_name)
 
-@router.delete("/{name}")
-def delete_script(name: str, qnum: Optional[int] = Query(None, description="If given, only remove the unit for this qnum; otherwise remove all units for the script")):
+
+@router.delete("/{name}", response_model=Dict[str, Any])
+def delete_script(name: str, qnum: Optional[int] = Query(None, description="If given, only remove the unit for this qnum; otherwise remove all units for the script")) -> Dict[str, Any]:
     """
     Delete a script and its associated resources.
     - If qnum is provided: stop/remove fw-script-<name>-q<qnum>.service (if present).
@@ -624,8 +668,8 @@ def delete_script(name: str, qnum: Optional[int] = Query(None, description="If g
     spath = script_path_for(name)
     rpath = requirements_path_for(name)
     vpath = venv_path_for(name)
-    file_removed = []
-    file_failed = []
+    file_removed: List[str] = []
+    file_failed: List[str] = []
 
     try:
         if os.path.exists(spath):
@@ -657,7 +701,7 @@ def delete_script(name: str, qnum: Optional[int] = Query(None, description="If g
         file_failed.append(vpath)
         errors.append(f"remove_venv {vpath}: {e}")
 
-    result = {
+    result: Dict[str, Any] = {
         "name": name,
         "units_removed": removed_units,
         "units_failed": failed_units,
@@ -668,10 +712,11 @@ def delete_script(name: str, qnum: Optional[int] = Query(None, description="If g
     logger.info("Delete script %s completed: removed_units=%d files_removed=%d errors=%d", name, len(removed_units), len(file_removed), len(errors))
     return result
 
+
 # ---------- Lifecycle helper ----------
-def register_lifecycle(app):
+def register_lifecycle(app) -> None:
     @app.on_event("shutdown")
-    def _shutdown_event():
+    def _shutdown_event() -> None:
         logger.info("Shutdown: stopping/removing manager-created units with prefix %s", UNIT_PREFIX)
         units = list_fw_units()
         for svc in units:
@@ -684,4 +729,4 @@ def register_lifecycle(app):
                     remove_unit(svc)
                 except Exception:
                     logger.exception("Failed to remove unit %s on shutdown", svc)
-        logger.info("Shutdown cleanup complete")
+        logger.info("Shutdown cleanup complete")