Onboard an Enterprise or MSSP Tenant

import json
import os
from pathlib import Path

import requests

BASE_URL = "https://fullhunt.io/api/v1/enterprise"
API_KEY = os.environ["FULLHUNT_API_KEY"]
TENANT_ID = os.environ["FULLHUNT_TENANT_ID"]
ORG_NAME = os.environ["FULLHUNT_ORG_NAME"]
DOMAINS = [v.strip() for v in os.environ.get("FULLHUNT_DOMAINS", "").split(",") if v.strip()]
IP_RANGES = [v.strip() for v in os.environ.get("FULLHUNT_IP_RANGES", "").split(",") if v.strip()]
AUTH_REF = os.environ["FULLHUNT_AUTHORIZATION_REFERENCE"]
STATE_PATH = Path(os.environ.get("FULLHUNT_ONBOARDING_STATE", f"tenant-{TENANT_ID}.json"))
HEADERS = {"X-API-KEY": API_KEY, "Content-Type": "application/json"}


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 request(session, method, path, **kwargs):
    response = session.request(
        method,
        f"{BASE_URL}{path}",
        headers=HEADERS,
        timeout=30,
        **kwargs,
    )
    response.raise_for_status()
    return response.json()


def resolve_organization(session, state):
    organizations = request(session, "GET", "/organizations")
    if state.get("organization_id"):
        if not any(org.get("id") == state["organization_id"] for org in organizations):
            raise RuntimeError("Stored organization is not available to this API key")
        return state["organization_id"]

    matches = [org for org in organizations if org.get("name") == ORG_NAME]
    if len(matches) > 1:
        raise RuntimeError("Multiple organizations have the requested name; use a stored ID")
    if matches:
        organization_id = matches[0]["id"]
    else:
        created = request(
            session,
            "POST",
            "/organizations",
            json={"name": ORG_NAME, "description": f"MSSP tenant {TENANT_ID}"},
        )
        organization_id = created["id"]

    state["organization_id"] = organization_id
    save_state(state)
    return organization_id


def add_scope(session, organization_id, state, path, field, values):
    completed = set(state.setdefault("registered_scope", []))
    for value in values:
        key = f"{field}:{value}"
        if key in completed:
            continue
        request(session, "POST", f"/organizations/{organization_id}/{path}", json={field: value})
        completed.add(key)
        state["registered_scope"] = sorted(completed)
        save_state(state)


def main():
    if not DOMAINS and not IP_RANGES:
        raise ValueError("Provide at least one authorized domain or public IP range")

    state = {
        "tenant_id": TENANT_ID,
        "organization_name": ORG_NAME,
        "authorization_reference": AUTH_REF,
    }
    if STATE_PATH.exists():
        state.update(json.loads(STATE_PATH.read_text(encoding="utf-8")))
        if state.get("tenant_id") != TENANT_ID:
            raise ValueError("State file belongs to another tenant")

    with requests.Session() as session:
        organization_id = resolve_organization(session, state)
        add_scope(session, organization_id, state, "domains", "domain", DOMAINS)
        add_scope(session, organization_id, state, "ip-ranges", "ip_range", IP_RANGES)

        entities = request(session, "GET", "/entities", params={"org": organization_id})
        alerts = request(session, "GET", "/alerts", params={"org": organization_id, "page": 1})
        vulnerabilities = request(
            session,
            "GET",
            "/vulnerabilities",
            params={"org": organization_id},
        )
        state["initial_baseline"] = {
            "entities": len(entities),
            "alerts_first_page": len(alerts),
            "vulnerabilities": len(vulnerabilities),
        }
        save_state(state)

    print(json.dumps(state, indent=2))


if __name__ == "__main__":
    main()