script improvements

backend/src/api/packet_scripting_api.py

@@ -1,14 +1,15 @@
 """
 NFQUEUE Python-Scripting API Router
 Endpoints:
- - POST /scripts                -> upload script (multipart): script file, optional requirements file, required 'name' form field
- - GET  /scripts                -> list scripts
- - GET  /scripts/{name}         -> download script (binary blob)
- - POST /scripts/{name}/enable  -> enable systemd service for script on given qnum
- - POST /scripts/{name}/disable -> disable service for script on qnum
- - DELETE /scripts/{name}       -> disable all or a specific qnum service(s) and delete script + venv + requirements
- - GET  /scripts/status         -> status of all fw-script units
- - GET  /scripts/{name}/status  -> status of units for that script
+ - POST /scripts                        -> upload script (multipart): script file, optional requirements file, required 'name' form field
+ - GET  /scripts                        -> list scripts + per-script unit mappings/status (combined)
+ - GET  /scripts/{name}                 -> download script (binary blob)
+ - GET  /scripts/{name}/requirements    -> download requirements file (binary blob) if present
+ - PUT  /scripts/{name}/requirements    -> upload/replace requirements file (multipart). ALWAYS runs pip install and returns pip output.
+ - DELETE /scripts/{name}/requirements -> delete only requirements file and remove venv (cleanup)
+ - POST /scripts/{name}/enable          -> enable systemd service for script on given qnum
+ - POST /scripts/{name}/disable         -> disable service for script on qnum
+ - DELETE /scripts/{name}               -> disable all or a specific qnum service(s) and delete script + venv + requirements
 """
 
 import os
@@ -21,7 +22,7 @@ import logging
 from typing import Optional, List, Dict, Any
 
 from fastapi import APIRouter, UploadFile, File, Form, HTTPException, Query
-from fastapi.responses import FileResponse
+from fastapi.responses import FileResponse, JSONResponse
 from pydantic import BaseModel, Field
 
 # ---------- Configuration ----------
@@ -364,6 +365,22 @@ class ScriptInfo(BaseModel):
     path: str
 
 
+class UnitMapping(BaseModel):
+    service: str
+    parsed: UnitParsed
+    active: bool = False
+
+
+class ScriptWithStatus(ScriptInfo):
+    """
+    Represents a script plus discovered unit mappings (if any).
+    - mappings: list of UnitMapping for that script
+    - requirements_exists: whether a requirements file exists on disk
+    """
+    mappings: List[UnitMapping] = []
+    requirements_exists: bool = False
+
+
 class ScriptUploadResponse(ScriptInfo):
     pip: Optional[Dict[str, str]] = None
 
@@ -375,12 +392,6 @@ class EnableRequest(BaseModel):
     enable_at_boot: Optional[bool] = False
 
 
-class UnitMapping(BaseModel):
-    service: str
-    parsed: UnitParsed
-    active: bool = False
-
-
 class StatusForNameResponse(BaseModel):
     name: str
     mappings: List[UnitMapping]
@@ -395,45 +406,47 @@ class OperationResult(BaseModel):
 
 
 # ---------- Endpoints ----------
-# Note: Place status endpoints before the dynamic GET /{name} route to avoid routing conflicts.
+# NOTE: combined status info into GET /scripts below (replaces separate /status & /{name}/status endpoints)
 
-@router.get("/status", response_model=Dict[str, UnitStatus])
-def status_all() -> Dict[str, UnitStatus]:
+@router.get("", response_model=List[ScriptWithStatus])
+def list_scripts_with_status() -> List[ScriptWithStatus]:
     """
-    Discover all fw-script units via systemctl and report parsed ExecStart + active state.
+    Return list of scripts plus per-script unit mappings/status.
+    This combines the former /scripts and /scripts/status endpoints so clients get everything in one call.
+    Each entry contains:
+      - name, path
+      - mappings: list of UnitMapping (service, parsed ExecStart, active flag)
+      - requirements_exists: boolean
     """
-    units = list_fw_units()
-    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] = UnitStatus(parsed=parsed, active=active)
-    logger.debug("Status queried: found %d units", len(results))
-    return results
+    out: List[ScriptWithStatus] = []
 
+    # build a list of units once for efficiency
+    all_units = list_fw_units()
+    logger.debug("list_scripts_with_status: discovered %d fw units", len(all_units))
 
-@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: List[UnitMapping] = []
-    prefix = f"{UNIT_PREFIX}-{name}-q"
-    for svc in units:
-        if svc.startswith(prefix):
-            parsed = parse_unit_execstart(svc)
-            try:
-                active = is_unit_active(svc)
-            except Exception:
-                active = False
-            matches.append(UnitMapping(service=svc, parsed=parsed, active=active))
-    logger.debug("Status for %s -> %d matches", name, len(matches))
-    return StatusForNameResponse(name=name, mappings=matches)
+    for fn in os.listdir(SCRIPT_DIR):
+        if not fn.endswith(".py"):
+            continue
+        name = fn.rsplit(".", 1)[0]
+        spath = os.path.join(SCRIPT_DIR, fn)
+
+        # find units matching this script
+        prefix = f"{UNIT_PREFIX}-{name}-q"
+        mappings: List[UnitMapping] = []
+        for svc in all_units:
+            if svc.startswith(prefix):
+                parsed = parse_unit_execstart(svc)
+                try:
+                    active = is_unit_active(svc)
+                except Exception:
+                    active = False
+                mappings.append(UnitMapping(service=svc, parsed=parsed, active=active))
+
+        req_exists = os.path.exists(requirements_path_for(name))
+        out.append(ScriptWithStatus(name=name, path=spath, mappings=mappings, requirements_exists=req_exists))
+
+    logger.debug("Listed %d scripts with status", len(out))
+    return out
 
 
 @router.post("", response_model=ScriptUploadResponse)
@@ -455,6 +468,10 @@ async def upload_script(
         logger.warning("Invalid name provided: %s", name)
         raise HTTPException(status_code=400, detail=str(e))
 
+    if not script.filename.endswith(".py"):
+        logger.warning("Upload rejected: script not .py (name=%s original=%s)", name, script.filename)
+        raise HTTPException(status_code=400, detail="only .py scripts allowed")
+
     spath = script_path_for(name)
     if os.path.exists(spath):
         logger.warning("Upload rejected: script with name already exists: %s", name)
@@ -482,6 +499,7 @@ async def upload_script(
                 fh.write(req_data)
             logger.info("Saved requirements for %s at %s", name, req_path)
             try:
+                # create venv (if needed) and run pip install; pip output is returned in response
                 create_venv(name)
                 venv_created = True
                 res = pip_install_requirements(name, req_path)
@@ -500,6 +518,7 @@ async def upload_script(
                         shutil.rmtree(venv_path_for(name), ignore_errors=True)
                 except Exception:
                     logger.exception("Cleanup after pip failure partially failed for %s", name)
+                # Return pip failure as 500 with message
                 raise HTTPException(status_code=500, detail=f"pip install failed: {err_msg}")
     except HTTPException:
         raise
@@ -528,18 +547,6 @@ async def upload_script(
     return ScriptUploadResponse(**resp)
 
 
-@router.get("", response_model=List[ScriptInfo])
-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(ScriptInfo(name=name, path=os.path.join(SCRIPT_DIR, fn)))
-    logger.debug("Listed %d scripts", len(out))
-    return out
-
-
 @router.get(
     "/{name}",
     response_class=FileResponse,
@@ -571,6 +578,140 @@ def download_script(name: str) -> FileResponse:
     return FileResponse(path, media_type="application/octet-stream", filename=f"{name}.py")
 
 
+@router.get(
+    "/{name}/requirements",
+    response_class=FileResponse,
+    responses={
+        200: {
+            "content": {"application/octet-stream": {"schema": {"type": "string", "format": "binary"}}},
+            "description": "requirements.txt file (binary).",
+        },
+        404: {"description": "Not found"},
+        400: {"description": "Invalid name"},
+    },
+)
+def download_requirements(name: str) -> FileResponse:
+    """
+    Download the stored requirements file for a script as binary blob.
+    """
+    try:
+        validate_name(name)
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+    path = requirements_path_for(name)
+    if not os.path.exists(path):
+        raise HTTPException(status_code=404, detail="requirements not found")
+    logger.info("Download requirements for %s", name)
+    return FileResponse(path, media_type="application/octet-stream", filename=f"{name}-requirements.txt")
+
+
+@router.put(
+    "/{name}/requirements",
+    responses={
+        200: {"description": "Requirements replaced and pip output returned"},
+        400: {"description": "Invalid name or bad request"},
+        500: {"description": "pip install failed or storage error"},
+    },
+)
+async def upload_requirements_install(name: str, requirements: UploadFile = File(...)) -> Dict[str, Any]:
+    """
+    Replace / upload the requirements file for a given script.
+    - requirements: multipart file upload (UploadFile) — keeps Swagger UI file input.
+    - This endpoint ALWAYS runs pip install -r <file> into the script's venv and returns pip stdout/stderr.
+      If pip install fails, a 500 error is returned with the pip failure message.
+    """
+    try:
+        validate_name(name)
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+
+    # ensure script exists (we don't allow attaching requirements to non-existing script)
+    spath = script_path_for(name)
+    if not os.path.exists(spath):
+        raise HTTPException(status_code=404, detail="script not found")
+
+    if not requirements:
+        raise HTTPException(status_code=400, detail="requirements file required")
+
+    req_path = requirements_path_for(name)
+    try:
+        data = await requirements.read()
+        with open(req_path, "wb") as fh:
+            fh.write(data)
+        logger.info("Saved requirements for %s at %s", name, req_path)
+    except Exception as e:
+        logger.exception("Failed to write requirements for %s: %s", name, e)
+        raise HTTPException(status_code=500, detail="failed to save requirements")
+
+    # Now ALWAYS install and return pip output (raise 500 on failure)
+    try:
+        create_venv(name)
+        res = pip_install_requirements(name, req_path)
+        logger.info("pip install completed for %s via requirements upload", name)
+        return {"pip": {"stdout": res.get("stdout", ""), "stderr": res.get("stderr", "")}}
+    except Exception as e:
+        logger.exception("pip install failed for %s: %s", name, e)
+        # keep the requirements file for inspection; return 500 with details
+        raise HTTPException(status_code=500, detail=f"pip install failed: {str(e)}")
+
+
+@router.delete(
+    "/{name}/requirements",
+    responses={
+        200: {"description": "requirements removed and venv cleaned up"},
+        404: {"description": "script or requirements not found"},
+        500: {"description": "cleanup error"},
+    },
+)
+def delete_requirements_and_cleanup(name: str) -> Dict[str, Any]:
+    """
+    Delete only the requirements file for a script and attempt to remove the script's venv directory.
+    Returns a summary of what was removed and any errors.
+    """
+    try:
+        validate_name(name)
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+
+    spath = script_path_for(name)
+    if not os.path.exists(spath):
+        raise HTTPException(status_code=404, detail="script not found")
+
+    rpath = requirements_path_for(name)
+    vpath = venv_path_for(name)
+
+    removed = {"requirements_removed": False, "venv_removed": False}
+    errors: List[str] = []
+
+    # remove requirements file
+    try:
+        if os.path.exists(rpath):
+            os.remove(rpath)
+            removed["requirements_removed"] = True
+            logger.info("Removed requirements file %s", rpath)
+        else:
+            logger.debug("No requirements file to remove for %s", name)
+    except Exception as e:
+        logger.exception("Failed removing requirements file %s: %s", rpath, e)
+        errors.append(f"remove_requirements {rpath}: {e}")
+
+    # remove venv directory
+    try:
+        if os.path.isdir(vpath):
+            shutil.rmtree(vpath, ignore_errors=False)
+            removed["venv_removed"] = True
+            logger.info("Removed venv directory %s", vpath)
+        else:
+            logger.debug("No venv directory to remove for %s", name)
+    except Exception as e:
+        logger.exception("Failed removing venv %s: %s", vpath, e)
+        errors.append(f"remove_venv {vpath}: {e}")
+
+    if errors:
+        return JSONResponse(status_code=500, content={"removed": removed, "errors": errors})
+    return {"removed": removed}
+
+
 @router.post("/{name}/enable", response_model=OperationResult)
 def enable_script(name: str, req: EnableRequest) -> OperationResult:
     try: