scipting improvements

backend/src/api/packet_scripting_api.py

@@ -1,40 +1,56 @@
-# script_router_stateless.py
+# script_router_with_venv.py
 """
-Stateless APIRouter that manages uploaded scripts via systemd only.
+APIRouter for uploading scripts (with optional requirements.txt),
-- Upload/list/download scripts (stored in SCRIPT_DIR)
+creating per-script venvs (when requirements provided), and managing
-- Enable script on qnum => write systemd unit fw-script-<sid>-q<qnum>.service and start it
+systemd services that run scripts with a queue number argument.
-- Disable script on qnum => stop service and remove unit file
+
-- Status endpoints discover active services by scanning /etc/systemd/system for fw-script-*.service
+Usage:
-No DB or run-dir used — systemd is the source of truth.
+    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:
+# ---------- Venv helpers ----------
-    """Write unit file and daemon-reload systemd."""
+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)
+        logger.exception("systemctl stop failed for %s", service_name)
-    unit_path = unit_path_for(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 {
+        return {"service": service_name, "exec_start": exec_start, "sid": sd["sid"], "script_path": sd["script"], "qnum": int(sd["qnum"]), "extra": sd.get("extra") or ""}
-            "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 ----------
+# ---------- End utilities ----------
+
+# ---------- Endpoints ----------
+
 @router.post("", response_model=ScriptInfo)
-async def upload_script(file: UploadFile = File(...)):
+async def upload_script(script: UploadFile = File(...), requirements: Optional[UploadFile] = File(None)):
-    if not file.filename.endswith(".py"):
+    """
-        logger.warning("Reject upload (not .py): %s", file.filename)
+    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)
+    logger.info("Uploaded script %s as sid=%s path=%s", script.filename, sid, path)
-    return {"id": sid, "name": file.filename, "path": 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)})
+        p = os.path.join(SCRIPT_DIR, fn)
-    logger.debug("Listed %d scripts", len(out))
+        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}")
+        write_unit(service_name, exec_start, description=f"FW script {sid} queue {qnum}", enable_at_boot=req.enable_at_boot)
-        # small pause
+        # 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)
+    except Exception as e:
-    return {"status": "ok", "sid": sid, "qnum": qnum, "service": service_name}
+        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)
+        # try stopping anyway in case only registered with systemd without unit file
-        # 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
+        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
+    try:
+        stop_unit(service_name)
+    except subprocess.CalledProcessError:
+        logger.exception("Failed stopping service %s", service_name)
+    try:
         remove_unit(service_name)
-    logger.info("Disabled script %s on qnum=%s (removed service %s)", sid, qnum, 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