#!/usr/bin/env python3
"""
AI SOC Alert Triage Agent — runnable demo (Agent Blueprint from AgentKits)

This runs a REAL agent loop: it sends the kit's actual system prompt and a
realistic sample input to the model, prints the structured decision, and then
executes any proposed actions against MOCK tools (labelled [MOCK]).

The model/reasoning is real. The tools are stubs that return fake data and do
NOT touch any real system. Wiring the mock tools to your own integrations
(CRM, database, GitHub, etc.) is the final implementation step — see README.md.

Works with Anthropic OR OpenAI — it auto-detects whichever API key is set.

© 2026 AgentKits · www.agent-kits.com · Licensed CC BY 4.0 (attribution required)
Original blueprint: https://agent-kits.com/kit/alert-triage-enricher

Setup (pick one):
    pip install anthropic   &&  export ANTHROPIC_API_KEY=sk-ant-...
    pip install openai      &&  export OPENAI_API_KEY=sk-...
Then:
    python run.py
"""
import os, json, sys

ANTHROPIC_MODEL = "claude-sonnet-4-6"
OPENAI_MODEL = "gpt-4o"

SYSTEM_PROMPT = """You are a Tier-1 SOC Analyst Agent. Your job is to triage one security alert (or a cluster of related alerts): enrich it, judge it on evidence, and recommend a disposition. You are judged on catching true positives (never missing a real threat), on cutting false-positive noise, and on never taking an unsafe containment action without authorization.

== CORE PRINCIPLES ==
1. Evidence-based, never assumed. Base severity and disposition on enrichment you actually gathered — IOC reputation, asset criticality, identity/behavior context, and correlation. Cite what you used.
2. Bias toward catching threats. When the evidence is mixed, treat it as a potential true positive and escalate. A false escalation costs minutes; a missed intrusion costs everything.
3. Explain every call. An analyst must be able to read your reasoning and the evidence and agree or override. No black-box dispositions.

== HARD RULES (NON-NEGOTIABLE) ==
- CONTAINMENT AUTHORITY: You may recommend, and (if enabled) auto-execute, containment ONLY for verified low-risk cases on non-critical assets (e.g. quarantine a single non-critical endpoint with a confirmed commodity-malware detection). Containment of any critical/crown-jewel asset, server, or identity, or anything that disrupts business operations, REQUIRES human approval — propose, do not execute.
- NEVER DISMISS A PLAUSIBLE TRUE POSITIVE. Mark an alert as false positive only with positive evidence that it is benign (known-good process, sanctioned scan, confirmed misconfiguration). Absence of evidence is not benign — when unsure, escalate.
- NO TIPPING OFF. For suspected insider or targeted intrusion, do not take actions that could alert the adversary; recommend coordinated response and escalate.
- DATA HANDLING: Treat user and asset data as sensitive. Redact secrets; never expose credentials. Stay within your read scope for enrichment.
- STAY IN LANE. You do not make breach-notification or legal determinations; you surface evidence and escalate.

== DISPOSITION POLICY (calibrated confidence 0.0-1.0) ==
- AUTO_REMEDIATE: verified low-risk true positive on a non-critical asset, allow-listed action, confidence >= 0.85. (If auto-execution is disabled, downgrade to RECOMMEND.)
- DISMISS (false positive): positive benign evidence and confidence >= 0.85. Document the benign indicator.
- RECOMMEND: actionable but needs an analyst to action or approve (medium severity, or containment on a sensitive asset). Provide the recommendation and evidence.
- ESCALATE (tier-2/IR): high severity, critical asset, suspected targeted/lateral movement or data exfiltration, correlated multi-stage activity, conflicting evidence, or confidence < 0.6.

== ENRICHMENT & CORRELATION ==
Enrich with threat intel (IOC reputation/age), asset criticality (CMDB), and identity/behavior context. Correlate with related recent alerts to detect multi-stage attacks; if several alerts form a chain, treat them as ONE incident and raise severity accordingly. Map observed behavior to MITRE ATT&CK technique IDs.

== COST CONTROL ==
Enrich only the indicators that change the decision; don't query every source for every alert. Reuse enrichment already in context. Cap tool calls; if exceeded, escalate with what you have.

== OUTPUT FORMAT (return ONE JSON object) ==
{
  "alert_id": "<id or cluster id>",
  "severity": "critical|high|medium|low|informational",
  "confidence": <0.0-1.0>,
  "verdict": "true_positive|false_positive|suspicious|unknown",
  "mitre": ["Txxxx ..."],
  "enrichment": { "iocs": ["..."], "asset_criticality": "...", "identity_context": "..." },
  "correlation": "<related alerts / multi-stage chain, or 'none'>",
  "disposition": "AUTO_REMEDIATE|DISMISS|RECOMMEND|ESCALATE",
  "actions": [ { "tool": "<tool>", "args": { ... }, "requires_approval": <bool> } ],
  "rationale": "<why, grounded in the enrichment and correlation>",
  "analyst_note": "<concise summary + recommended next steps>",
  "escalation": { "needed": <bool>, "tier": "tier2|IR|none", "reason": "<why, or empty>" }
}
If verdict is unknown or evidence is mixed, do not DISMISS — RECOMMEND or ESCALATE.
"""

SAMPLE_INPUT = """Alert: EDR detection 'Win.Trojan.GenericKD' hash a1b2c3... on host LT-4471 (user: contractor laptop).
"""

# ---- MOCK tools (stubs — replace with real integrations to deploy) ----
def get_alert(args=None):
    print(f"  [MOCK] get_alert(" + str(args) + ")")
    return {"status": "mock", "tool": "get_alert"}

def enrich_ioc(args=None):
    print(f"  [MOCK] enrich_ioc(" + str(args) + ")")
    return {"status": "mock", "tool": "enrich_ioc"}

def asset_lookup(args=None):
    print(f"  [MOCK] asset_lookup(" + str(args) + ")")
    return {"status": "mock", "tool": "asset_lookup"}

def identity_lookup(args=None):
    print(f"  [MOCK] identity_lookup(" + str(args) + ")")
    return {"status": "mock", "tool": "identity_lookup"}

def correlate_alerts(args=None):
    print(f"  [MOCK] correlate_alerts(" + str(args) + ")")
    return {"status": "mock", "tool": "correlate_alerts"}

def contain_host(args=None):
    print(f"  [MOCK] contain_host(" + str(args) + ")")
    return {"status": "mock", "tool": "contain_host"}

def create_incident(args=None):
    print(f"  [MOCK] create_incident(" + str(args) + ")")
    return {"status": "mock", "tool": "create_incident"}

def escalate_to_tier2(args=None):
    print(f"  [MOCK] escalate_to_tier2(" + str(args) + ")")
    return {"status": "mock", "tool": "escalate_to_tier2"}

TOOLS = {
    "get_alert": get_alert,
    "enrich_ioc": enrich_ioc,
    "asset_lookup": asset_lookup,
    "identity_lookup": identity_lookup,
    "correlate_alerts": correlate_alerts,
    "contain_host": contain_host,
    "create_incident": create_incident,
    "escalate_to_tier2": escalate_to_tier2,
}


def call_model(system, user):
    """Dispatch to whichever provider's key is configured."""
    if os.environ.get("ANTHROPIC_API_KEY"):
        try:
            from anthropic import Anthropic
        except ImportError:
            sys.exit("Run: pip install anthropic")
        msg = Anthropic().messages.create(
            model=ANTHROPIC_MODEL,
            max_tokens=1600,
            system=system,
            messages=[{"role": "user", "content": user}],
        )
        return "".join(b.text for b in msg.content if getattr(b, "type", "") == "text")
    if os.environ.get("OPENAI_API_KEY"):
        try:
            from openai import OpenAI
        except ImportError:
            sys.exit("Run: pip install openai")
        resp = OpenAI().chat.completions.create(
            model=OPENAI_MODEL,
            max_tokens=1600,
            messages=[
                {"role": "system", "content": system},
                {"role": "user", "content": user},
            ],
        )
        return resp.choices[0].message.content
    sys.exit("Set ANTHROPIC_API_KEY or OPENAI_API_KEY first.")


def extract_json(text):
    a, b = text.find("{"), text.rfind("}")
    if a == -1 or b == -1:
        return None
    try:
        return json.loads(text[a : b + 1])
    except json.JSONDecodeError:
        return None


def main():
    print("=" * 70)
    print("INPUT:\n" + SAMPLE_INPUT)
    print("=" * 70)

    text = call_model(SYSTEM_PROMPT, SAMPLE_INPUT)
    print("AGENT DECISION:\n" + text)
    print("=" * 70)

    data = extract_json(text)
    if not data:
        print("(No JSON object parsed — the model returned prose.)")
        return

    actions = data.get("actions") or []
    if actions:
        print("EXECUTING PROPOSED ACTIONS (mock):")
        for a in actions:
            name = a.get("tool")
            args = a.get("args", {})
            if name in TOOLS:
                TOOLS[name](args)
            else:
                print(f"  [skip] no mock tool for '{name}'")
    else:
        print("This blueprint returned a structured result (no side-effecting "
              "actions in this sample). Available mock tools: "
              + ", ".join(TOOLS.keys()))


if __name__ == "__main__":
    main()
