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

"""
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
 - 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
"""

import os
import sys
import re
import shutil
import subprocess
import time
import logging
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, Field

# ---------- Configuration ----------
SCRIPT_DIR = "/srv/fw-scripts"
VENV_BASE = "/srv/fw-scripts/venvs"
UNIT_DIR = "/etc/systemd/system"   # retained for writing new units, but discovery uses systemctl
UNIT_PREFIX = "fw-script"

# ensure dirs exist
os.makedirs(SCRIPT_DIR, exist_ok=True)
os.makedirs(VENV_BASE, exist_ok=True)

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

# ---------- Router ----------
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")
    if not _NAME_RE.match(name):
        raise ValueError("invalid name; allowed characters: letters, digits, dot, underscore, hyphen")
    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-.
    Uses `systemctl list-units` to be canonical.
    """
    units: List[str] = []
    try:
        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():
            line = line.strip()
            if not line:
                continue
            cols = line.split()
            unit = cols[0]
            if unit.startswith(UNIT_PREFIX + "-") and unit.endswith(".service"):
                units.append(unit[:-8])  # strip ".service"
    except Exception:
        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).
    Returns None if unknown / empty.
    """
    unit = _systemctl_unit_name(service_name)
    try:
        p = subprocess.run(["systemctl", "show", "-p", "FragmentPath", "--value", unit],
                           capture_output=True, text=True, check=False, timeout=2)
        frag = (p.stdout or "").strip()
        if not frag:
            return None
        return frag
    except Exception:
        logger.exception("systemctl show FragmentPath failed for %s", service_name)
        return None


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 UnitParsed instance.
    """
    unit = service_name if service_name.endswith(".service") else service_name + ".service"

    try:
        p = subprocess.run(
            ["systemctl", "--no-pager", "show", "-p", "ExecStart", "--value", unit],
            capture_output=True,
            text=True,
            check=False,
            timeout=3,
        )
        raw = (p.stdout or "").strip()
    except Exception:
        logger.exception("systemctl show ExecStart failed for %s", service_name)
        return UnitParsed(service=service_name)

    if not raw:
        return UnitParsed(service=service_name)

    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: List[str] = []
        for a in argv_entries:
            v = a.strip()
            # strip surrounding quotes if present
            if (v.startswith('"') and v.endswith('"')) or (v.startswith("'") and v.endswith("'")):
                v = v[1:-1]
            if v:
                argv_parts.append(v)
        if argv_parts:
            exec_start = " ".join(argv_parts)

    # Look for /srv/fw-scripts/<name>.py directly
    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 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: Optional[int] = None
    extra: Optional[str] = None
    if after:
        tokens = after.split()
        for i, t in enumerate(tokens):
            try:
                val = int(t)
                if 0 <= val <= 65535:
                    qnum = val
                    extra = " ".join(tokens[i+1:]) if i + 1 < len(tokens) else ""
                    break
            except ValueError:
                continue

    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:
    unit = _systemctl_unit_name(service_name)
    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]
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)
    subprocess.run(["systemctl", "daemon-reload"], check=True)
    logger.info("Wrote unit %s", unit_path)
    if enable_at_boot:
        try:
            subprocess.run(["systemctl", "enable", _systemctl_unit_name(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 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),
    and daemon-reload. All via systemctl queries (no fallback to scanning /etc).
    """
    unit = _systemctl_unit_name(service_name)
    try:
        subprocess.run(["systemctl", "stop", unit], check=False)
    except Exception:
        logger.exception("Failed stopping %s", unit)
    try:
        subprocess.run(["systemctl", "disable", unit], check=False)
    except Exception:
        pass

    frag = get_unit_fragment_path(service_name)
    if frag:
        try:
            if os.path.exists(frag):
                os.remove(frag)
                logger.info("Removed unit file at %s for %s", frag, service_name)
        except Exception:
            logger.exception("Failed to remove unit file %s", frag)
    else:
        logger.debug("No FragmentPath for %s; nothing to delete on-disk", service_name)

    # ensure systemd reloads units
    try:
        subprocess.run(["systemctl", "daemon-reload"], check=True)
    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)
    if os.path.exists(venv_dir):
        logger.debug("Venv already exists for %s", name)
        return venv_dir
    os.makedirs(os.path.dirname(venv_dir), exist_ok=True)
    logger.info("Creating venv for %s at %s", name, venv_dir)
    try:
        subprocess.run([sys.executable, "-m", "venv", venv_dir], check=True, timeout=timeout, capture_output=True, text=True)
    except subprocess.CalledProcessError as e:
        logger.exception("venv creation failed for %s: %s", name, e.stderr if hasattr(e, "stderr") else str(e))
        raise RuntimeError("venv creation failed: " + (e.stderr or str(e)))
    except subprocess.TimeoutExpired:
        logger.exception("venv creation timed out for %s", name)
        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")
    try:
        subprocess.run([pip_path, "install", "--upgrade", "pip"], check=True, capture_output=True, text=True, timeout=300)
    except subprocess.CalledProcessError as e:
        logger.warning("pip upgrade warning for %s: %s", name, getattr(e, "stderr", str(e)))
    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 succeeded for %s", name)
        return {"stdout": p.stdout or "", "stderr": p.stderr or ""}
    except subprocess.CalledProcessError as e:
        out = {"stdout": getattr(e, "stdout", "") or "", "stderr": getattr(e, "stderr", "") or str(e)}
        logger.error("pip install failed for %s: %s", name, out["stderr"][:4000])
        raise RuntimeError(jsonify_cmd_output(out))
    except subprocess.TimeoutExpired:
        logger.error("pip install timed out for %s", name)
        raise RuntimeError("pip install timed out")


# ---------- Models ----------
class ScriptInfo(BaseModel):
    name: str
    path: str


class ScriptUploadResponse(ScriptInfo):
    pip: Optional[Dict[str, str]] = None


class EnableRequest(BaseModel):
    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", 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, 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


@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)


@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.
    """
    # validate name
    try:
        validate_name(name)
    except ValueError as e:
        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)
        raise HTTPException(status_code=409, detail="script with that name already exists")

    data = await script.read()
    try:
        with open(spath, "wb") as fh:
            fh.write(data)
        os.chmod(spath, 0o700)
        logger.info("Saved script for name=%s at %s", name, spath)
    except Exception as e:
        logger.exception("Failed to write script file for %s: %s", name, e)
        raise HTTPException(status_code=500, detail="failed to save script")

    req_path: Optional[str] = None
    pip_output: Optional[Dict[str, str]] = None
    venv_created = False

    try:
        if requirements is not None:
            req_data = await requirements.read()
            req_path = requirements_path_for(name)
            with open(req_path, "wb") as fh:
                fh.write(req_data)
            logger.info("Saved requirements for %s at %s", name, req_path)
            try:
                create_venv(name)
                venv_created = True
                res = pip_install_requirements(name, req_path)
                pip_output = {"stdout": res.get("stdout", ""), "stderr": res.get("stderr", "")}
                logger.info("pip install completed for %s", name)
            except Exception as pip_exc:
                err_msg = str(pip_exc)
                logger.error("pip install error for %s: %s", name, err_msg[:2000])
                # cleanup: remove script, req file, partial venv
                try:
                    if os.path.exists(spath):
                        os.remove(spath)
                    if req_path and os.path.exists(req_path):
                        os.remove(req_path)
                    if venv_created and os.path.isdir(venv_path_for(name)):
                        shutil.rmtree(venv_path_for(name), ignore_errors=True)
                except Exception:
                    logger.exception("Cleanup after pip failure partially failed for %s", name)
                raise HTTPException(status_code=500, detail=f"pip install failed: {err_msg}")
    except HTTPException:
        raise
    except Exception as e:
        logger.exception("Unexpected error during upload for %s: %s", name, e)
        try:
            if os.path.exists(spath):
                os.remove(spath)
        except Exception:
            pass
        try:
            if req_path and os.path.exists(req_path):
                os.remove(req_path)
        except Exception:
            pass
        try:
            if os.path.isdir(venv_path_for(name)):
                shutil.rmtree(venv_path_for(name), ignore_errors=True)
        except Exception:
            pass
        raise HTTPException(status_code=500, detail="internal error during upload")

    resp = {"name": name, "path": spath}
    if pip_output is not None:
        resp["pip"] = pip_output
    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}")
def download_script(name: str):
    # dynamic route - placed after /status and /{name}/status
    try:
        validate_name(name)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    path = script_path_for(name)
    if not os.path.exists(path):
        logger.warning("Download requested for missing script %s", name)
        raise HTTPException(status_code=404, detail="not found")
    logger.info("Download script %s", name)
    return FileResponse(path, media_type="text/x-python", filename=f"{name}.py")


@router.post("/{name}/enable", response_model=OperationResult)
def enable_script(name: str, req: EnableRequest) -> OperationResult:
    try:
        validate_name(name)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))

    script_path = script_path_for(name)
    if not os.path.exists(script_path):
        raise HTTPException(status_code=404, detail="script not found")

    qnum = req.qnum
    service_name = req.service_name or make_service_name(name, qnum)
    python_path = venv_python_for(name)
    exec_start = f"{python_path} {script_path} {qnum}"
    if req.extra_args:
        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 or False)
        time.sleep(0.05)
        start_unit(service_name)
    except subprocess.CalledProcessError as e:
        logger.exception("Failed to start service %s for %s: %s", service_name, name, e)
        try:
            remove_unit(service_name)
        except Exception:
            pass
        raise HTTPException(status_code=500, detail=f"systemd start failed: {e}")
    except Exception as e:
        logger.exception("Unknown error starting service %s for %s: %s", service_name, name, e)
        try:
            remove_unit(service_name)
        except Exception:
            pass
        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 OperationResult(status="ok", name=name, qnum=qnum, service=service_name, python=python_path)


@router.post("/{name}/disable", response_model=OperationResult)
def disable_script(name: str, qnum: int) -> OperationResult:
    try:
        validate_name(name)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    service_name = make_service_name(name, qnum)
    # attempt stop + remove via systemctl-based remove_unit
    try:
        # If unit exists according to systemctl, remove it; otherwise still try stopping
        units = list_fw_units()
        if service_name in units:
            remove_unit(service_name)
            logger.info("Disabled and removed unit %s", service_name)
        else:
            # attempt stop anyway
            subprocess.run(["systemctl", "stop", _systemctl_unit_name(service_name)], check=False)
            logger.info("Tried stopping unit %s (unit file not present)", service_name)
    except Exception as e:
        logger.exception("Failed to disable unit %s: %s", service_name, e)
        raise HTTPException(status_code=500, detail=str(e))
    return OperationResult(status="ok", name=name, qnum=qnum, service=service_name)


@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).
    - If qnum is not provided: stop/remove all fw-script-<name>-q*.service units found.
    - Remove script file, requirements file, and venv directory.
    Returns JSON summarizing performed actions and any errors.
    """
    try:
        validate_name(name)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))

    removed_units: List[str] = []
    failed_units: List[str] = []
    errors: List[str] = []

    # determine units to handle
    if qnum is not None:
        svc = make_service_name(name, qnum)
        units_to_handle = [svc]
    else:
        all_units = list_fw_units()
        prefix = f"{UNIT_PREFIX}-{name}-q"
        units_to_handle = [u for u in all_units if u.startswith(prefix)]

    # stop and remove units
    for svc in units_to_handle:
        try:
            try:
                subprocess.run(["systemctl", "stop", _systemctl_unit_name(svc)], check=False)
            except Exception:
                pass
            # remove unit (systemd-based)
            remove_unit(svc)
            removed_units.append(svc)
        except Exception as e:
            logger.exception("Error removing unit %s: %s", svc, e)
            failed_units.append(svc)
            errors.append(f"remove_unit {svc}: {e}")

    # remove files
    spath = script_path_for(name)
    rpath = requirements_path_for(name)
    vpath = venv_path_for(name)
    file_removed: List[str] = []
    file_failed: List[str] = []

    try:
        if os.path.exists(spath):
            os.remove(spath)
            file_removed.append(spath)
            logger.info("Removed script file %s", spath)
    except Exception as e:
        logger.exception("Failed removing script file %s: %s", spath, e)
        file_failed.append(spath)
        errors.append(f"remove_script {spath}: {e}")

    try:
        if os.path.exists(rpath):
            os.remove(rpath)
            file_removed.append(rpath)
            logger.info("Removed requirements file %s", rpath)
    except Exception as e:
        logger.exception("Failed removing requirements file %s: %s", rpath, e)
        file_failed.append(rpath)
        errors.append(f"remove_requirements {rpath}: {e}")

    try:
        if os.path.isdir(vpath):
            shutil.rmtree(vpath, ignore_errors=False)
            file_removed.append(vpath)
            logger.info("Removed venv directory %s", vpath)
    except Exception as e:
        logger.exception("Failed removing venv %s: %s", vpath, e)
        file_failed.append(vpath)
        errors.append(f"remove_venv {vpath}: {e}")

    result: Dict[str, Any] = {
        "name": name,
        "units_removed": removed_units,
        "units_failed": failed_units,
        "files_removed": file_removed,
        "files_failed": file_failed,
        "errors": errors,
    }
    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) -> None:
    @app.on_event("shutdown")
    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:
            if svc.startswith(UNIT_PREFIX + "-"):
                try:
                    subprocess.run(["systemctl", "stop", _systemctl_unit_name(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")