#!/usr/bin/env python3
"""Datto EDR (Infocyte HUNT) REST API client for the datto-edr skill.

Datto EDR == rebranded Infocyte HUNT. The API is a per-tenant LoopBack REST
service at https://<instance>.infocyte.com/api/<Model>. This client talks to the
live ACG tenant (azcomp4587).

Auth (IMPORTANT): a single long-lived 64-char token passed in the RAW
`Authorization` header — NOT "Bearer <token>", no prefix. (Verified live
2026-06-25.) Token + instance load from the SOPS vault; env overrides exist for
testing.

Transport: prefers httpx if installed, else stdlib urllib (no hard dependency).

Read methods are always live. Mutating methods (scan / isolate / deploy) return
the raw upstream result; the CLI caller is responsible for gating them behind
--confirm.
"""
from __future__ import annotations

import json
import os
import subprocess
import sys
import urllib.error
import urllib.parse
import urllib.request
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any, Optional

try:
    import httpx  # type: ignore

    _HAS_HTTPX = True
except ImportError:  # pragma: no cover - depends on environment
    _HAS_HTTPX = False

# Cap upstream error bodies surfaced in exceptions. `raw` can call arbitrary
# paths whose responses may reflect other data - bound the blast radius.
ERROR_BODY_MAX_CHARS = 500

# --- constants ----------------------------------------------------------------
DEFAULT_INSTANCE = "azcomp4587"
DATTO_EDR_BASE_URL = os.environ.get(
    "DATTO_EDR_BASE_URL", f"https://{DEFAULT_INSTANCE}.infocyte.com/api"
)
TIMEOUT_SECONDS = 60.0
CONNECT_TIMEOUT_SECONDS = 10.0

VAULT_ENTRY = "msp-tools/datto-edr.sops.yaml"
VAULT_FIELD = "credentials.api_token"

SKILL_DIR = Path(__file__).resolve().parent.parent

# Mutating-path guard for `raw`: refuse these without --confirm at the call site.
# Scans, isolation responses, deletes, approvals all run through POST/DELETE.


class DattoEDRError(RuntimeError):
    """Raised for transport or API errors."""


# --- credential loading -------------------------------------------------------
def _resolve_claudetools_root() -> Path:
    """Resolve the ClaudeTools repo root: env var, then identity.json, then derived."""
    derived_root = SKILL_DIR.parent.parent.parent  # skills/datto-edr -> repo root
    env_root = os.environ.get("CLAUDETOOLS_ROOT")
    if env_root:
        return Path(env_root)
    identity_path = derived_root / ".claude" / "identity.json"
    if identity_path.exists():
        try:
            data = json.loads(identity_path.read_text(encoding="utf-8"))
            root = data.get("claudetools_root")
            if root:
                return Path(root)
        except (json.JSONDecodeError, OSError):
            pass
    return derived_root


def load_api_token() -> str:
    """Load the Datto EDR API token: DATTO_EDR_TOKEN env override, else SOPS vault."""
    env_key = os.environ.get("DATTO_EDR_TOKEN")
    if env_key:
        return env_key.strip()

    root = _resolve_claudetools_root()
    vault_script = root / ".claude" / "scripts" / "vault.sh"
    if not vault_script.exists():
        raise DattoEDRError(
            f"Cannot load API token: vault wrapper not found at {vault_script} "
            "and DATTO_EDR_TOKEN is not set."
        )
    try:
        completed = subprocess.run(
            ["bash", str(vault_script), "get-field", VAULT_ENTRY, VAULT_FIELD],
            capture_output=True,
            text=True,
            timeout=60,
        )
    except FileNotFoundError as exc:
        raise DattoEDRError(
            "Cannot load API token: 'bash' not found on PATH. Install Git Bash or "
            "set DATTO_EDR_TOKEN."
        ) from exc
    except subprocess.TimeoutExpired as exc:
        raise DattoEDRError("Cannot load API token: vault call timed out.") from exc

    if completed.returncode != 0:
        raise DattoEDRError(
            "Cannot load API token from vault "
            f"(exit {completed.returncode}): {completed.stderr.strip()}"
        )
    token = completed.stdout.strip()
    if not token:
        raise DattoEDRError("Vault returned an empty API token.")
    return token


# --- client -------------------------------------------------------------------
class DattoEDRClient:
    def __init__(
        self,
        api_token: Optional[str] = None,
        api_base_url: str = DATTO_EDR_BASE_URL,
        timeout: float = TIMEOUT_SECONDS,
        connect_timeout: float = CONNECT_TIMEOUT_SECONDS,
    ):
        self.api_base_url = api_base_url.rstrip("/")
        self._api_token = api_token
        self.timeout = timeout
        self.connect_timeout = connect_timeout

    @property
    def api_token(self) -> str:
        if not self._api_token:
            self._api_token = load_api_token()
        return self._api_token

    # -- core transport --------------------------------------------------------
    def _request(self, method: str, path: str, filt: Optional[dict] = None,
                 body: Optional[dict] = None) -> Any:
        """One REST call. `filt` -> ?filter=<json> (LoopBack). Returns parsed JSON."""
        url = f"{self.api_base_url}/{path.lstrip('/')}"
        if filt is not None:
            qs = urllib.parse.urlencode({"filter": json.dumps(filt)})
            url = f"{url}?{qs}"
        data = json.dumps(body).encode("utf-8") if body is not None else None
        result = self._send(method, url, data)
        if isinstance(result, dict) and result.get("error"):
            err = result["error"]
            msg = err.get("message") if isinstance(err, dict) else err
            raise DattoEDRError(f"Datto EDR API error [{method} {path}]: {msg}")
        return result

    def _send(self, method: str, url: str, data: Optional[bytes]) -> Any:
        # Auth is the RAW token in the Authorization header (no "Bearer").
        headers = {
            "Authorization": self.api_token,
            "Content-Type": "application/json",
            "Accept": "application/json",
        }
        if _HAS_HTTPX:
            try:
                timeout = httpx.Timeout(self.timeout, connect=self.connect_timeout)
                with httpx.Client(timeout=timeout) as client:
                    resp = client.request(method, url, content=data, headers=headers)
                resp.raise_for_status()
                return resp.json() if resp.content else None
            except httpx.TimeoutException as exc:
                raise DattoEDRError(f"Datto EDR request timed out: {exc}") from exc
            except httpx.HTTPStatusError as exc:
                detail = (exc.response.text or "")[:ERROR_BODY_MAX_CHARS]
                raise DattoEDRError(
                    f"Datto EDR HTTP {exc.response.status_code}: {detail}"
                ) from exc
            except httpx.HTTPError as exc:
                raise DattoEDRError(f"Datto EDR request failed: {exc}") from exc

        # stdlib fallback
        req = urllib.request.Request(url, data=data, method=method, headers=headers)
        try:
            with urllib.request.urlopen(req, timeout=self.timeout) as resp:
                raw = resp.read()
            return json.loads(raw.decode("utf-8")) if raw else None
        except urllib.error.HTTPError as exc:
            detail = exc.read().decode("utf-8", errors="replace")[:ERROR_BODY_MAX_CHARS]
            raise DattoEDRError(f"Datto EDR HTTP {exc.code}: {detail}") from exc
        except urllib.error.URLError as exc:
            raise DattoEDRError(f"Datto EDR request failed: {exc}") from exc

    # convenience wrappers
    def _get(self, path: str, filt: Optional[dict] = None) -> Any:
        return self._request("GET", path, filt=filt)

    def _count(self, model: str, where: Optional[dict] = None) -> int:
        path = f"{model}/count"
        if where:
            qs = urllib.parse.urlencode({"where": json.dumps(where)})
            path = f"{path}?{qs}"
        res = self._request("GET", path)
        if isinstance(res, dict):
            return int(res.get("count", 0))
        return 0

    # ======================================================================
    # READ METHODS (always live)
    # ======================================================================
    def get_status(self) -> dict:
        """Tenant-wide rollup: counts of agents/alerts/orgs/targets."""
        return {
            "instance": self.api_base_url,
            "organizations": self._count("Organizations"),
            "targetGroups": self._count("Targets"),
            "agents": self._count("Agents"),
            "agentsActive": self._count("Agents", {"active": True}),
            "agentsIsolated": self._count("Agents", {"isolated": True}),
            "alerts": self._count("Alerts"),
        }

    def list_organizations(self, limit: int = 200) -> list[dict]:
        """List Organizations (= client tenants/companies). VERIFIED LIVE."""
        return self._get("Organizations", {
            "limit": limit, "order": "name ASC",
            "fields": {"id": True, "name": True, "agentCount": True,
                       "alertCount": True, "locationCount": True},
        }) or []

    def list_locations(self, org_id: Optional[str] = None, limit: int = 500) -> list[dict]:
        """List Locations (= client sites). VERIFIED LIVE. This is the per-client
        grouping that agents belong to (Agent.locationId -> Location.id), and a
        Location carries organizationId. Org -> Locations -> Agents is the
        inventory hierarchy."""
        filt: dict = {"limit": limit, "order": "name ASC",
                      "fields": {"id": True, "name": True, "organizationId": True,
                                 "agentCount": True, "activeAgentCount": True,
                                 "alertCount": True, "lastScannedOn": True}}
        if org_id:
            filt["where"] = {"organizationId": org_id}
        return self._get("Locations", filt) or []

    def list_targets(self, org_id: Optional[str] = None, limit: int = 500) -> list[dict]:
        """List Targets (= scan groups). VERIFIED LIVE. A grouping concept that
        belongs to an Organization via organizationId (often aliases a Location
        id). NOTE: scanning is per-AGENT via Agents/scan - the legacy
        targets/{id}/scan route is DEAD (404) on this tenant."""
        filt: dict = {"limit": limit, "order": "name ASC",
                      "fields": {"id": True, "name": True, "organizationId": True,
                                 "agentCount": True, "activeAgentCount": True,
                                 "alertCount": True, "lastScannedOn": True}}
        if org_id:
            filt["where"] = {"organizationId": org_id}
        return self._get("Targets", filt) or []

    def list_agents(self, location_id: Optional[str] = None,
                    location_ids: Optional[list[str]] = None,
                    limit: int = 500) -> list[dict]:
        """List Agents (endpoints). VERIFIED LIVE. Filter by locationId (a single
        site) or location_ids (an org's sites, via LoopBack `inq`)."""
        filt: dict = {"limit": limit, "order": "hostname ASC"}
        where: dict = {}
        if location_id:
            where["locationId"] = location_id
        elif location_ids:
            where["locationId"] = {"inq": location_ids}
        if where:
            filt["where"] = where
        return self._get("Agents", filt) or []

    def get_agent(self, agent_id: str) -> dict:
        """Agent detail by id. VERIFIED LIVE."""
        return self._get(f"Agents/{agent_id}") or {}

    def list_alerts(self, org_id: Optional[str] = None,
                    target_group_id: Optional[str] = None,
                    severity: Optional[int] = None, days: Optional[int] = None,
                    limit: int = 200) -> list[dict]:
        """List Alerts (detections). VERIFIED LIVE.
        Alerts carry organizationId/organizationName + targetGroupId/Name inline."""
        where: dict = {}
        if org_id:
            where["organizationId"] = org_id
        if target_group_id:
            where["targetGroupId"] = target_group_id
        if severity is not None:
            where["severity"] = severity
        if days is not None:
            since = (datetime.now(timezone.utc) - timedelta(days=days)).isoformat()
            where["createdOn"] = {"gt": since}
        filt: dict = {"limit": limit, "order": "createdOn DESC"}
        if where:
            filt["where"] = where
        return self._get("Alerts", filt) or []

    def get_alert(self, alert_id: str) -> dict:
        """Alert detail by id. VERIFIED LIVE."""
        return self._get(f"Alerts/{alert_id}") or {}

    def list_agent_keys(self) -> Any:
        """List agent registration keys (/agentKeys). Used to build deploy commands."""
        return self._get("agentKeys") or []

    def sweep(self, org_id: Optional[str] = None) -> dict:
        """Security posture rollup per client: agent health + AV + isolation +
        recent detection counts. Built from live reads (never cached)."""
        orgs = self.list_organizations()
        if org_id:
            orgs = [o for o in orgs if o.get("id") == org_id]
        since = (datetime.now(timezone.utc) - timedelta(days=7)).isoformat()
        out = []
        for o in orgs:
            oid = o.get("id")
            # _count (not list+len) so a >500-alert client isn't silently capped
            # and mis-sorted, and we don't transfer 500 full records per org.
            recent = self._count("Alerts", {"organizationId": oid,
                                            "createdOn": {"gt": since}})
            out.append({
                "organizationId": oid,
                "organization": o.get("name"),
                "agents": o.get("agentCount", 0),
                "alertsTotal": o.get("alertCount", 0),
                "alertsLast7d": recent,
            })
        out.sort(key=lambda r: (-r["alertsLast7d"], (r["organization"] or "").lower()))
        return {"clients": out}

    # ======================================================================
    # MUTATING METHODS (caller MUST gate behind --confirm)
    # ----------------------------------------------------------------------
    # The Infocyte-module scan routes (targets/{id}/scan, targets/scan, scans)
    # are DEAD (404) on Datto EDR. The live mechanism, read verbatim from the
    # console JS bundle, is:
    #   POST agents/scan {where:<loopback where>, options:<scan opts>, taskName}
    # where `where` SELECTS the agents. An ABSENT/empty `where` scans the ENTIRE
    # tenant (the 156-host footgun) - so scan_agents REFUSES an empty id list.
    # Response extensions (isolate/kill/quarantine) ride the same call via
    # options.extensions. (Verified live 2026-06-25.)
    # ======================================================================
    # Full EDR forensic collection - the default so a bare `scan` is a real sweep
    # (an empty options body risks a thin scan that returns a false all-clear on a
    # compromised host). Pass options={} explicitly to override with a lean scan.
    DEFAULT_SCAN_OPTIONS = {
        "process": True, "module": True, "driver": True, "memory": True,
        "account": True, "artifact": True, "autostart": True, "application": True,
        "installed": True, "events": True,
    }

    def scan_agents(self, agent_ids: list[str],
                    options: Optional[dict] = None,
                    task_name: str = "Scan - EDR") -> Any:
        """Trigger an EDR scan on specific agents (POST agents/scan).

        Targets by Agent.id via the AND-wrapped LoopBack where the console uses:
          {"where":{"and":[{"id":[ids]}]}, "options":{...}, "taskName":"Scan - EDR"}.
        A bare {"id":{"inq":[...]}} returns HTTP 412 "column reference id is
        ambiguous" (the scan query joins tables) - the and-wrap disambiguates.
        REFUSES an empty agent_ids list (an absent where scans the whole tenant).
        `options` defaults to DEFAULT_SCAN_OPTIONS (full forensic); pass {} for lean.
        STATE-CHANGING - gate behind --confirm at the call site."""
        ids = [i for i in (agent_ids or []) if i]
        if not ids:
            raise DattoEDRError(
                "scan_agents refused: empty agent list. An absent `where` scans "
                "the ENTIRE tenant (156-host footgun) - pass explicit agent id(s)."
            )
        body = {
            "where": {"and": [{"id": ids}]},
            "options": self.DEFAULT_SCAN_OPTIONS if options is None else options,
            "taskName": task_name,
        }
        return self._request("POST", "Agents/scan", body=body)

    def list_extensions(self, limit: int = 200) -> Any:
        """List agent extensions (response + collection). Read.
        Response extensions (e.g. Host Isolation) are invoked via scan options."""
        return self._get("Extensions", {"limit": limit}) or []

    def run_extension_on_agents(self, agent_ids: list[str], extension_id: str,
                                task_name: str = "Response") -> Any:
        """Invoke a response extension (isolate/kill/quarantine) on specific
        agents - rides POST agents/scan with the extension in options.extensions.
        Same empty-list guard as scan_agents. Confirm the extension id via
        list_extensions() (e.g. 'Host Isolation [Win/Linux]').
        STATE-CHANGING - gate behind --confirm at the call site."""
        ids = [i for i in (agent_ids or []) if i]
        if not ids:
            raise DattoEDRError(
                "run_extension_on_agents refused: empty agent list (would target "
                "the whole tenant). Pass explicit agent id(s)."
            )
        if not extension_id:
            raise DattoEDRError("run_extension_on_agents requires an extension_id.")
        body = {
            "where": {"and": [{"id": ids}]},  # AND-wrapped (see scan_agents)
            "options": {"extensions": [{"id": extension_id}]},
            "taskName": task_name,
        }
        return self._request("POST", "Agents/scan", body=body)

    # -- task status / control --------------------------------------------------
    def list_tasks(self, type_: Optional[str] = None, limit: int = 20) -> list[dict]:
        """List userTasks (scan/analysis jobs), newest first. Read."""
        filt: dict = {"limit": limit, "order": "createdOn DESC"}
        if type_:
            filt["where"] = {"type": type_}
        return self._get("userTasks", filt) or []

    def get_task(self, task_id: str) -> dict:
        """Get one userTask (scan job) detail. Read."""
        return self._get(f"userTasks/{task_id}") or {}

    def cancel_task(self, task_id: str) -> Any:
        """Cancel a running userTask, e.g. a scan (POST userTasks/{id}/cancel ->
        204). STATE-CHANGING - gate behind --confirm at the call site."""
        return self._request("POST", f"userTasks/{task_id}/cancel")

    # -- provisioning -----------------------------------------------------------
    def create_group(self, name: str, organization_id: str) -> Any:
        """Create an EDR target group under an org (POST Targets).
        STATE-CHANGING - gate behind --confirm at the call site."""
        return self._request("POST", "Targets",
                             body={"name": name, "organizationId": organization_id})

    def mint_registration_key(self, target_id: str, key: str) -> Any:
        """Mint an agent registration key tied to a target group (POST agentKeys).
        The key string `id` is CALLER-SUPPLIED (10-char); an absent id 500s.
        STATE-CHANGING - gate behind --confirm at the call site."""
        return self._request("POST", "agentKeys",
                             body={"id": key, "targetId": target_id})

    # ======================================================================
    # POWER TOOL
    # ======================================================================
    def raw(self, method: str, path: str, filt: Optional[dict] = None,
            body: Optional[dict] = None) -> Any:
        """Call any endpoint directly. Caller gates mutating methods."""
        return self._request(method.upper(), path, filt=filt, body=body)


def main() -> int:
    """Minimal self-check: load token (no network call)."""
    try:
        client = DattoEDRClient()
        _ = client.api_token  # triggers vault load
        print("[OK] API token loaded; transport =",
              "httpx" if _HAS_HTTPX else "urllib")
        print("[INFO] instance =", client.api_base_url)
        return 0
    except DattoEDRError as exc:
        print(f"[ERROR] {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())
