Orchestrate an OEM Scan Lifecycle

import json
import os
import random
import time
from pathlib import Path

import requests

BASE_URL = "https://fullhunt.io/api/v1/oem/on-demand-scan"
API_KEY = os.environ["FULLHUNT_API_KEY"]
SCAN_TYPE = os.environ["FULLHUNT_SCAN_TYPE"]
TARGET = os.environ["FULLHUNT_SCAN_TARGET"]
TENANT_ID = os.environ["FULLHUNT_TENANT_ID"]
JOB_ID = os.environ["FULLHUNT_JOB_ID"]
STATE_PATH = Path(os.environ.get("FULLHUNT_SCAN_STATE", "oem-scan-state.json"))
TIMEOUT_SECONDS = int(os.environ.get("FULLHUNT_SCAN_TIMEOUT", "1800"))
FORCE_NEW_SCAN = os.environ.get("FULLHUNT_FORCE_NEW_SCAN") == "1"
HEADERS = {"X-API-KEY": API_KEY, "Content-Type": "application/json"}
TAGS = {"tenant_id": TENANT_ID, "job_id": JOB_ID, "workflow": "scan-lifecycle"}


def post(session, path, body, attempts=5):
    for attempt in range(attempts):
        response = session.post(
            f"{BASE_URL}{path}",
            headers=HEADERS,
            json=body,
            timeout=30,
        )
        if response.status_code == 429 or response.status_code >= 500:
            retry_after = response.headers.get("Retry-After")
            delay = float(retry_after) if retry_after else min(2**attempt, 30)
            time.sleep(delay + random.uniform(0, 1))
            continue
        response.raise_for_status()
        return response.json()
    raise RuntimeError(f"{path or '/'} did not recover after {attempts} attempts")


def load_state():
    if not STATE_PATH.exists():
        return None
    state = json.loads(STATE_PATH.read_text(encoding="utf-8"))
    if state.get("target") != TARGET or state.get("type") != SCAN_TYPE:
        raise ValueError("State file belongs to a different target or scan type")
    tags = state.get("query_tags") or {}
    if tags.get("tenant_id") != TENANT_ID or tags.get("job_id") != JOB_ID:
        raise ValueError("State file belongs to a different tenant or job")
    return state


def save_state(state):
    temporary = STATE_PATH.with_suffix(STATE_PATH.suffix + ".tmp")
    temporary.write_text(json.dumps(state, indent=2) + "\n", encoding="utf-8")
    temporary.replace(STATE_PATH)


def main():
    if SCAN_TYPE not in {"domain", "host", "ip_range"}:
        raise ValueError("FULLHUNT_SCAN_TYPE must be domain, host, or ip_range")

    with requests.Session() as session:
        state = load_state()
        if state and state.get("status") == "scan_completed" and not FORCE_NEW_SCAN:
            print(json.dumps(state, indent=2))
            return
        if not state or FORCE_NEW_SCAN:
            queued = post(
                session,
                "",
                {"type": SCAN_TYPE, "target": TARGET, "query_tags": TAGS},
            )["response"]
            state = {
                "target": TARGET,
                "type": SCAN_TYPE,
                "scan_id": queued["scan_id"],
                "queued_at": queued.get("timestamp"),
                "status": "queued",
                "query_tags": TAGS,
            }
            save_state(state)

        deadline = time.monotonic() + TIMEOUT_SECONDS
        while time.monotonic() < deadline:
            payload = post(
                session,
                "/scan-status",
                {"scan_id": state["scan_id"], "query_tags": TAGS},
            )["response"]
            scan_info = payload["scan_info"]
            status = str(scan_info.get("status", "unknown")).lower()
            state.update({"status": status, "scan_info": scan_info})
            save_state(state)

            if status == "scan_completed":
                print(json.dumps(state, indent=2))
                return
            if "fail" in status or status == "error":
                raise RuntimeError(f"Scan reached failure state: {status}")

            time.sleep(30)

    raise TimeoutError("Scan did not reach a terminal state before timeout")


if __name__ == "__main__":
    main()