app.py

# Flask API to run Appollo scans as Kubernetes Jobs (and legacy GET /run_appollo).

import base64
import os
import re
import sys
import time
import traceback
import random
import string
from flask import Flask, request, jsonify

try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    pass

from k8s import send_job_to_k8s, get_job_status, get_job_pod_logs
from public_app_notes import register_public_app_notes_routes

app = Flask(__name__)
register_public_app_notes_routes(app)

# In-memory store for pod callbacks (status + logs). Key = job_id, value = dict.
# Used when dashboard cannot read K8s logs; pods POST here when they start/finish.
# Also populated from Kubernetes Job API when the pod callback never arrives (see job_status).
_JOB_CALLBACKS = {}
_MAX_LOG_BYTES = 512 * 1024  # 512KB per job
# job_name -> unix time when appollo-api created the Job (for not_found after TTL cleanup)
_SUBMITTED_JOBS = {}
_MAX_SUBMITTED_TRACK = 2000

# K8s Job defaults (override via env)
SCAN_NAMESPACE = os.environ.get("SCAN_NAMESPACE", "appollo-scans")
SCAN_IMAGE = os.environ.get("SCAN_IMAGE", "appollo:latest")
CONFIG_MAP_NAME = os.environ.get("SCAN_CONFIG_MAP", "appollo-configs")
CONFIG_SECRET_NAME = os.environ.get("SCAN_CONFIG_SECRET", "appollo-secrets")
PVC_CLAIM_NAME = os.environ.get("SCAN_PVC_CLAIM", "pvc-rwm")
IMAGE_PULL_SECRET = os.environ.get("SCAN_IMAGE_PULL_SECRET", "registry-secret")
SERVICE_ACCOUNT = os.environ.get("SCAN_SERVICE_ACCOUNT", "appollo-sa")
ENV_PATH_IN_CONTAINER = "/app/config/.env"
# URL the scan pod uses to call back. Must be FQDN so it resolves from job namespace (e.g. appollo-scans).
# Appollo API may run in a different namespace (e.g. appollo-scans); use APPOLLO_API_NAMESPACE or set APPOLLO_API_INTERNAL_URL.
APPOLLO_API_NAMESPACE = os.environ.get("APPOLLO_API_NAMESPACE", "appollo-scans")
_CALLBACK_DEFAULT = f"http://appollo-api.{APPOLLO_API_NAMESPACE}.svc.cluster.local"
APPOLLO_API_INTERNAL_URL = os.environ.get("APPOLLO_API_INTERNAL_URL", _CALLBACK_DEFAULT).rstrip("/")
# After Job is removed from cluster (TTL), infer "completed" if still polling (seconds since submit).
# 900s left the dashboard stuck on "running" for 15+ minutes; 120s is enough for API lag + short GC.
JOB_NOT_FOUND_INFER_DONE_AFTER_SEC = int(
    os.environ.get("JOB_NOT_FOUND_INFER_DONE_AFTER_SEC", "120")
)
# How long finished Jobs stay in etcd before TTL deletes them (must match or exceed Job spec)
SCAN_JOB_TTL_AFTER_FINISH = int(os.environ.get("SCAN_JOB_TTL_AFTER_FINISH", "86400"))

# Dashboard scan IDs -> appollo.py CLI flags (must match dashboard)
SCAN_FLAGS = {
    "update-inventory":       "-U",
    "ssl-checker":            "-sc",
    "port-scan":              "-ps",
    "firewall-port-scan":     "-fs",
    "exposed-services-scan":  "-es",
    "ip-exposure-scan":       "-ipe",
    "external-probe-scan":    "-ep",
    "probe-all-scan":         "-pa",
    "cloud-functions-scan":   "-cf",
    "cloud-storage-scan":     "-cs",
    "cert-transparency-scan": "-ct",
    "wayback-scan":           "-ws",
    "tech-scan":              "-ts",
    "dir-scan":               "-ds",
    "dast-scan":              "-dast",
    "subdomain-scan":         "-sub",
    "email-security-scan":    "-em",
    "nuclei-scan":            "-ns",
    "dangling-dns":           "-dd",
    "aws-scan":               "-as",
    "vendor-scan":            "-vs",
}
ALLOWED_SCANS = set(SCAN_FLAGS.keys()) | {"complete-scan"}


def random_string(length=5):
    return "".join(random.choices(string.ascii_lowercase + string.digits, k=length))


def _safe_job_name(s):
    return re.sub(r"[^a-z0-9-]", "-", s.lower())[:32]


def build_appollo_command(target, scans):
    """
    Build: python3 src/appollo.py -e /app/config/.env [ -t target | -A ] [ flags ].
    - No target => use -A (full inventory). With target => -t <target>.
    - complete-scan => -A + all scan flags (all types). Otherwise selected flags only.
    - vendor-scan => uses -V <slug> instead of -t <target> (slug passed as target).
    """
    parts = ["python3", "src/appollo.py", "-e", ENV_PATH_IN_CONTAINER]
    is_vendor_scan = "vendor-scan" in scans
    if is_vendor_scan and target and target.strip():
        # Vendor slug is passed as target; map to -V flag
        parts.extend(["-V", target.strip()])
    elif target and target.strip():
        parts.extend(["-t", target.strip()])
    else:
        parts.append("-A")  # full inventory when no target
    if "complete-scan" in scans:
        for flag in SCAN_FLAGS.values():
            parts.append(flag)
    else:
        for s in scans:
            if s in SCAN_FLAGS:
                parts.append(SCAN_FLAGS[s])
    return " ".join(parts)


def _wrap_command_with_callback(job_name, appollo_cmd):
    """Wrap appollo command: on exit, POST completion to appollo-api (Python + curl fallback)."""
    appollo_escaped = appollo_cmd.replace("'", "'\"'\"'")
    delim = "CBSCRIPT"
    # Python first; if it fails (no urllib/ssl), curl/wget notifies appollo-api.
    script_end = """import urllib.request, json, os
url = os.environ.get("DASHBOARD_CALLBACK_URL", "").strip()
jid = os.environ.get("JOB_ID", "").strip()
st = os.environ.get("STATUS", "completed")
ex = os.environ.get("EXIT", "0")
if url and jid:
    payload = json.dumps({"jobId": jid, "status": st, "message": "Exit " + ex}).encode()
    u = url.rstrip("/") + "/callback/scan"
    for attempt in range(3):
        try:
            r = urllib.request.Request(u, data=payload, headers={"Content-Type": "application/json"}, method="POST")
            urllib.request.urlopen(r, timeout=20)
            raise SystemExit(0)
        except SystemExit:
            raise
        except Exception as e:
            print("[appollo-callback] python attempt %s failed: %s" % (attempt + 1, e), flush=True)
            if attempt < 2:
                import time
                time.sleep(2 * (attempt + 1))
    raise RuntimeError("python callback failed")
"""
    # shell: run python; on failure try curl then wget (typical scan images include one of these)
    shell_cb = (
        "python3 <<'" + delim + "'\n" + script_end + delim + "\n"
        "PY=$?; "
        "if [ \"$PY\" -ne 0 ] && [ -n \"$DASHBOARD_CALLBACK_URL\" ] && [ -n \"$JOB_ID\" ]; then "
        "CB=\"${DASHBOARD_CALLBACK_URL%/}/callback/scan\"; "
        "BODY=\"{\\\"jobId\\\":\\\"$JOB_ID\\\",\\\"status\\\":\\\"$STATUS\\\",\\\"message\\\":\\\"exit $EXIT\\\"}\"; "
        "(command -v curl >/dev/null 2>&1 && curl -sS -X POST -H 'Content-Type: application/json' "
        "-d \"$BODY\" \"$CB\" --connect-timeout 10 --max-time 45 && echo '[appollo-callback] curl ok') || "
        "echo '[appollo-callback] curl failed' >&2; "
        "fi"
    )
    return (
        "(" + appollo_escaped + ") 2>&1 | tee /tmp/scan.log; EXIT=$?; "
        "STATUS=completed; [ \"$EXIT\" -ne 0 ] && STATUS=failed; export EXIT STATUS JOB_ID; "
        + shell_cb + "; "
        "exit $EXIT"
    )


def create_scan_job(target, scans):
    """Create a Job spec that runs appollo in K8s (command form, like prod Pod spec)."""
    if not scans:
        raise ValueError("scans list is empty")
    scan_slug = "complete" if "complete-scan" in scans else _safe_job_name("-".join(scans[:3]))
    ts = int(time.time())
    job_name = f"appollo-scan-{scan_slug}-{ts}-{random_string()}"
    cmd = build_appollo_command(target, scans)
    wrapped_cmd = _wrap_command_with_callback(job_name, cmd)
    full_command = ["sh", "-c", wrapped_cmd]
    container_env = [
        {"name": "JOB_ID", "value": job_name},
        {"name": "DASHBOARD_CALLBACK_URL", "value": APPOLLO_API_INTERNAL_URL},
    ]
    job_body = {
        "apiVersion": "batch/v1",
        "kind": "Job",
        "metadata": {"name": job_name},
        "spec": {
            "ttlSecondsAfterFinished": SCAN_JOB_TTL_AFTER_FINISH,
            "template": {
                "metadata": {"labels": {"app": "appollo"}},
                "spec": {
                    "containers": [
                        {
                            "name": "appollo",
                            "image": SCAN_IMAGE,
                            "command": full_command,
                            "env": container_env,
                            "volumeMounts": [
                                {"name": "config-volume-1", "mountPath": "/app/config"},
                                {"name": "pvc-install", "mountPath": "/etc/config"},
                            ],
                            "imagePullPolicy": "Always",
                        }
                    ],
                    "volumes": [
                        {
                            "name": "config-volume-1",
                            "projected": {
                                "sources": [
                                    {"configMap": {"name": CONFIG_MAP_NAME, "optional": True}},
                                    {"secret": {"name": CONFIG_SECRET_NAME, "optional": True}},
                                ],
                                "defaultMode": 420,
                            },
                        },
                        {
                            "name": "pvc-install",
                            "persistentVolumeClaim": {"claimName": PVC_CLAIM_NAME},
                        },
                    ],
                    "imagePullSecrets": [{"name": IMAGE_PULL_SECRET}] if (IMAGE_PULL_SECRET or "").strip() else [],
                    "restartPolicy": "Never",
                    "serviceAccountName": SERVICE_ACCOUNT,
                },
            },
        },
    }
    return job_name, job_body


def create_job(target, job_args):
    """Legacy: create job with args list (for /run_appollo)."""
    job_name = f"appollo-{_safe_job_name(target or 'scan')}-{random_string()}"
    job_json = {
        "apiVersion": "batch/v1",
        "kind": "Job",
        "metadata": {"name": job_name},
        "spec": {
            "ttlSecondsAfterFinished": 60,
            "template": {
                "spec": {
                    "containers": [
                        {
                            "name": "scanjob",
                            "image": SCAN_IMAGE,
                            "imagePullPolicy": "Always",
                            "args": job_args,
                            "volumeMounts": [
                                {"name": "config-volume-1", "mountPath": "/app/config"},
                                {"name": "pvc-install", "mountPath": "/etc/config"},
                            ],
                        }
                    ],
                    "volumes": [
                        {"name": "config-volume-1", "configMap": {"name": CONFIG_MAP_NAME}},
                        {"name": "pvc-install", "persistentVolumeClaim": {"claimName": PVC_CLAIM_NAME}},
                    ],
                    "imagePullSecrets": [{"name": IMAGE_PULL_SECRET}],
                    "restartPolicy": "Never",
                }
            },
        },
    }
    return job_json


# ---- Routes ----


@app.route("/run_scan", methods=["POST"])
def run_scan():
    """
    Accept dashboard payload: { "target?: string", "scans": string[] }.
    Create a K8s Job that runs appollo with those scans. Returns job name for status polling.
    """
    if not request.is_json:
        return jsonify({"error": "Content-Type must be application/json"}), 400
    body = request.get_json() or {}
    scans = body.get("scans")
    if not isinstance(scans, list) or not scans:
        return jsonify({"error": "scans array is required and must be non-empty"}), 400
    target = body.get("target") if isinstance(body.get("target"), str) else None
    valid = [s for s in scans if s in ALLOWED_SCANS]
    if not valid:
        return jsonify({"error": "No valid scan types"}), 400
    try:
        job_name, job_body = create_scan_job(target, valid)
        res = send_job_to_k8s(job_body, namespace=SCAN_NAMESPACE)
        _register_submitted_job(res.metadata.name)
        return jsonify({
            "success": True,
            "jobId": res.metadata.name,
            "message": "Scan job submitted to Kubernetes.",
        }), 202
    except Exception as e:
        app.logger.exception("Failed to create scan job")
        print("[run_scan] ERROR:", str(e), file=sys.stderr)
        traceback.print_exc(file=sys.stderr)
        return jsonify({"error": "Failed to create scan job. Please try again later."}), 500


def _valid_job_id(name):
    """Allow only safe job names (K8s job names: lowercase, digits, hyphens)."""
    if not name or len(name) > 64:
        return False
    return bool(re.match(r"^[a-z0-9][a-z0-9.-]*$", name))


def _register_submitted_job(job_name):
    """Track jobs we created so not_found (TTL) can be interpreted as finished."""
    while len(_SUBMITTED_JOBS) >= _MAX_SUBMITTED_TRACK:
        oldest = min(_SUBMITTED_JOBS.items(), key=lambda x: x[1])[0]
        _SUBMITTED_JOBS.pop(oldest, None)
    _SUBMITTED_JOBS[job_name] = time.time()


def _merge_terminal_callback(job_name, status, message):
    """Store completion from K8s (or inference); do not overwrite pod-reported terminal state."""
    if status not in ("completed", "failed"):
        return
    if job_name not in _JOB_CALLBACKS:
        _JOB_CALLBACKS[job_name] = {"status": None, "message": None, "logs": None, "updated_at": None}
    cur = _JOB_CALLBACKS[job_name].get("status")
    if cur in ("completed", "failed"):
        return
    _JOB_CALLBACKS[job_name]["status"] = status
    _JOB_CALLBACKS[job_name]["message"] = (message or "").strip() or None
    _JOB_CALLBACKS[job_name]["updated_at"] = time.time()


def _enrich_status_response(job_name, out):
    """
    Always expose callback-style fields when Kubernetes already knows the outcome.
    Pod POST to /callback/scan is optional; many clusters block pod→appollo-api — K8s API is authoritative.
    """
    st = out.get("status")
    if st == "completed":
        out["callback_status"] = "completed"
        out["completion_source"] = "kubernetes"
        _merge_terminal_callback(
            job_name,
            "completed",
            "Job succeeded (reported by Kubernetes API).",
        )
    elif st == "failed":
        out["callback_status"] = "failed"
        out["completion_source"] = "kubernetes"
        _merge_terminal_callback(
            job_name,
            "failed",
            out.get("message") or "Job failed (reported by Kubernetes API).",
        )
    elif st == "not_found":
        submitted_at = _SUBMITTED_JOBS.get(job_name)
        if submitted_at and (time.time() - submitted_at) >= JOB_NOT_FOUND_INFER_DONE_AFTER_SEC:
            out["status"] = "completed"
            out["callback_status"] = "completed"
            out["completion_source"] = "inferred_ttl"
            out["message"] = (
                "Job no longer in the cluster (likely finished and garbage-collected). "
                "Treated as completed for dashboard status."
            )
            _merge_terminal_callback(job_name, "completed", out["message"])
        else:
            cb = _JOB_CALLBACKS.get(job_name)
            if cb and cb.get("status") in ("completed", "failed"):
                out["callback_status"] = cb["status"]
                out["callback_message"] = cb.get("message")
                out["status"] = cb["status"]
    elif st == "running":
        cb = _JOB_CALLBACKS.get(job_name)
        if cb and cb.get("status") in ("completed", "failed"):
            out["callback_status"] = cb["status"]
            out["completion_source"] = "pod_callback"
            out["callback_message"] = cb.get("message")
        elif cb and cb.get("status") == "started":
            out["callback_status"] = "started"
            out["completion_source"] = "pod_callback"
    cb = _JOB_CALLBACKS.get(job_name)
    if cb:
        if out.get("callback_status") is None and cb.get("status"):
            out["callback_status"] = cb["status"]
            out["callback_message"] = cb.get("message")
        if cb.get("logs") is not None:
            out["logs"] = cb["logs"]
    return out


def _logs_nonempty_for_fetch(out):
    """Safe check whether we should still pull pod logs (avoids .strip() on non-strings → 500)."""
    raw = out.get("logs") if isinstance(out, dict) else None
    if raw is None:
        return False
    try:
        if isinstance(raw, bytes):
            raw = raw.decode("utf-8", errors="replace")
        return bool(str(raw).strip())
    except Exception:
        return False


def _coerce_out_logs_to_utf8_str(out):
    """Ensure out['logs'] is a UTF-8 str and capped for JSON (Flask jsonify must not throw)."""
    if not isinstance(out, dict) or "logs" not in out or out["logs"] is None:
        return
    logs = out["logs"]
    if isinstance(logs, bytes):
        out["logs"] = logs.decode("utf-8", errors="replace")
    elif not isinstance(logs, str):
        out["logs"] = str(logs)
    if len(out["logs"]) > _MAX_LOG_BYTES:
        out["logs"] = out["logs"][-_MAX_LOG_BYTES:]


@app.route("/callback/scan", methods=["POST"])
def callback_scan():
    """
    Pod callback: report status and optional logs. Body: jobId, status, message?, logsBase64?.
    Status: started | completed | failed. Logs are stored and returned from GET /job/<id>/status and /job/<id>/logs.
    """
    if not request.is_json:
        return jsonify({"error": "Content-Type must be application/json"}), 400
    body = request.get_json() or {}
    job_id = (body.get("jobId") or "").strip()
    if not _valid_job_id(job_id):
        return jsonify({"error": "Invalid or missing jobId"}), 400
    status = (body.get("status") or "").strip().lower()
    if status not in ("started", "completed", "failed"):
        return jsonify({"error": "status must be one of: started, completed, failed"}), 400
    message = (body.get("message") or "").strip() or None
    logs_b64 = body.get("logsBase64")
    logs_text = None
    if isinstance(logs_b64, str) and logs_b64:
        try:
            raw = base64.b64decode(logs_b64, validate=True)
            logs_text = raw.decode("utf-8", errors="replace")
            if len(logs_text) > _MAX_LOG_BYTES:
                logs_text = logs_text[-_MAX_LOG_BYTES:]
        except Exception as e:
            app.logger.warning("Callback logsBase64 decode failed: %s", e)
    if job_id not in _JOB_CALLBACKS:
        _JOB_CALLBACKS[job_id] = {"status": None, "message": None, "logs": None, "updated_at": None}
    _JOB_CALLBACKS[job_id]["status"] = status
    _JOB_CALLBACKS[job_id]["message"] = message
    if logs_text is not None:
        _JOB_CALLBACKS[job_id]["logs"] = logs_text
    _JOB_CALLBACKS[job_id]["updated_at"] = time.time()
    app.logger.info(
        "scan callback received jobId=%s status=%s", job_id, status
    )
    return jsonify({"ok": True}), 200


@app.route("/job/<job_name>/status", methods=["GET"])
def job_status(job_name):
    """
    K8s Job status + callback fields. For valid job ids we always return HTTP 200 with JSON
    the dashboard can parse — never 500 (avoids 'Scan service error (500)' on polls).
    """
    if not job_name or "/" in job_name or not _valid_job_id(job_name):
        return jsonify({"error": "Invalid job name"}), 400
    try:
        try:
            out = get_job_status(SCAN_NAMESPACE, job_name)
        except Exception:
            app.logger.exception("get_job_status raised unexpectedly job=%s", job_name)
            out = {
                "status": "running",
                "message": "Scan service hit an internal error while querying Kubernetes.",
            }
        if not isinstance(out, dict) or not out.get("status"):
            out = {
                "status": "running",
                "message": "Invalid status payload from Kubernetes layer.",
            }
        _pre_enrich = dict(out)
        try:
            out = _enrich_status_response(job_name, out)
        except Exception:
            app.logger.exception("_enrich_status_response failed job=%s", job_name)
            out = _pre_enrich
            if not (out.get("message") or "").strip():
                out["message"] = (
                    "Could not merge callback metadata; status below is from Kubernetes only."
                )
        if not isinstance(out, dict) or not out.get("status"):
            out = {"status": "running", "message": "Status response was invalid after enrichment."}

        cb = _JOB_CALLBACKS.get(job_name)
        if cb and cb.get("message") and "callback_message" not in out:
            out["callback_message"] = cb.get("message")

        st = out.get("status")
        cb_st = out.get("callback_status")
        terminal = st in ("completed", "failed") or cb_st in ("completed", "failed")
        if terminal and not _logs_nonempty_for_fetch(out):
            try:
                pod_logs = get_job_pod_logs(SCAN_NAMESPACE, job_name)
                if pod_logs:
                    out["logs"] = pod_logs
                    if job_name not in _JOB_CALLBACKS:
                        _JOB_CALLBACKS[job_name] = {
                            "status": None,
                            "message": None,
                            "logs": None,
                            "updated_at": None,
                        }
                    _JOB_CALLBACKS[job_name]["logs"] = pod_logs
            except Exception as e:
                app.logger.warning("Could not fetch pod logs for %s: %s", job_name, e)

        _coerce_out_logs_to_utf8_str(out)
        return jsonify(out), 200
    except Exception:
        app.logger.exception("job_status fatal job=%s", job_name)
        return jsonify(
            {
                "status": "running",
                "message": "Scan service could not complete this status request; the next poll will retry.",
            }
        ), 200


@app.route("/job/<job_name>/logs", methods=["GET"])
def job_logs(job_name):
    """Return logs from callback cache or Kubernetes pod logs."""
    if not job_name or "/" in job_name or not _valid_job_id(job_name):
        return jsonify({"error": "Invalid job name"}), 400
    cb = _JOB_CALLBACKS.get(job_name)
    if cb and cb.get("logs"):
        return jsonify({"jobId": job_name, "logs": cb["logs"]}), 200
    try:
        pod_logs = get_job_pod_logs(SCAN_NAMESPACE, job_name)
    except Exception as e:
        app.logger.warning("job_logs K8s fetch failed %s: %s", job_name, e)
        pod_logs = None
    if pod_logs:
        if job_name not in _JOB_CALLBACKS:
            _JOB_CALLBACKS[job_name] = {
                "status": None,
                "message": None,
                "logs": None,
                "updated_at": None,
            }
        _JOB_CALLBACKS[job_name]["logs"] = pod_logs
        return jsonify({"jobId": job_name, "logs": pod_logs}), 200
    return jsonify({"error": "Logs not available for this job"}), 404


@app.route("/run_appollo", methods=["GET"])
def run_docker_command():
    params = request.args
    if not params or "target" not in params:
        return "Missing target parameter", 400
    target = params["target"]
    args = ["--target", target]
    scan_options = {
        "dnsscan": "--update-dns",
        "dirscan": "--dir-scan",
        "portscan": "--port-scan",
        "nucleiscan": "--nuclei-scan",
        "techscan": "--tech-scan",
        "realtimescan": "--realtime-scan",
        "waybackscan": "--wayback-scan",
    }
    for key, flag in scan_options.items():
        if params.get(key) == "true":
            args.append(flag)
    if len(args) == 2:
        return "No valid scans specified", 400
    job_scan = create_job(target, args)
    response = send_job_to_k8s(job_scan, namespace=SCAN_NAMESPACE)
    return jsonify({"message": "Job created successfully", "job": response.metadata.name}), 200


if __name__ == "__main__":
    port = int(os.environ.get("PORT", 1337))
    app.run(host="0.0.0.0", port=port, debug=False)