"""soc_agent.py, M9b: the SAME agent idea, built with a FRAMEWORK (LangGraph).

In 9a you wrote the tool loop by hand. A framework runs that loop for you, so you focus
on the tools and the task. This is the headline project: a SOC (Security Operations)
Level-1/2 analyst assistant that triages a security alert.

EDUCATIONAL, SYNTHETIC DATA ONLY. The threat-intel and logs are made up (see tools.py).
   Never point an agent like this at real systems, real logs, or real intel feeds without
   authorization. This assistant only investigates and summarizes, it takes no real action.

What it does (L1: enrich + summarize, L2: correlate):
  - lookup_ioc, reputation of an IP/domain/hash in the (synthetic) feed
  - search_logs, find related activity in the (synthetic) sample logs
  - it then correlates and summarizes, and remembers the conversation (memory).

Setup:  pip install langgraph langchain-anthropic   (Python 3.10-3.12; see install guide)
Run (venv active, key in .env, from this folder):
    python soc_agent.py
"""

import os
from dotenv import load_dotenv
from langchain_anthropic import ChatAnthropic
from langchain_core.tools import tool
from langgraph.prebuilt import create_react_agent
from langgraph.checkpoint.memory import MemorySaver

# Reuse the verified, synthetic-data functions from tools.py as the agent's tools.
from tools import lookup_ioc as _lookup_ioc, search_logs as _search_logs

load_dotenv()


@tool
def lookup_ioc(indicator: str) -> str:
    """Look up the reputation of an indicator (IP address, domain, or file hash)
    in the threat-intel feed. Use this to enrich indicators found in an alert."""
    return _lookup_ioc(indicator)


@tool
def search_logs(term: str) -> str:
    """Search the security logs for lines containing a term, such as an IP address
    or a username. Use this to find activity related to an alert."""
    return _search_logs(term)


SYSTEM = (
    "You are a SOC (Security Operations Center) Level-1/2 analyst assistant working with "
    "SAMPLE, synthetic data. Triage alerts methodically: enrich any indicators with "
    "lookup_ioc, find related activity with search_logs, then correlate the findings and "
    "give a short summary plus a suggested next step. Cite what each tool returned. You "
    "only investigate and summarize, you never take real-world action."
)

# create_react_agent runs the tool loop for you. We pass an explicit Claude model object
# (it reads ANTHROPIC_API_KEY from the environment). checkpointer=MemorySaver() gives the
# agent conversation memory, keyed by thread_id at invoke time.
model = ChatAnthropic(model="claude-opus-4-8", max_tokens=1024)
agent = create_react_agent(
    model,
    tools=[lookup_ioc, search_logs],
    prompt=SYSTEM,
    checkpointer=MemorySaver(),
)


if __name__ == "__main__":
    config = {"configurable": {"thread_id": "case-001"}}   # one investigation "thread"
    print("SOC assistant (SYNTHETIC data). Describe an alert or ask a follow-up. 'quit' to exit.")
    print("Try: 'Alert: repeated failed logins for admin from 185.220.101.45, triage it.'\n")
    while True:
        q = input("Analyst: ")
        if q.strip().lower() in {"quit", "exit"}:
            break
        result = agent.invoke({"messages": [{"role": "user", "content": q}]}, config)
        print("Assistant:", result["messages"][-1].content, "\n")