netforensicai 0.3.0__py3-none-any.whl

netforensicai/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

netforensicai/agents/__init__.py

@@ -0,0 +1,27 @@
+"""The investigation-team agents (see docs/design/agent-team.md).
+
+Each agent is one role - a mission, a scoped subset of the read-only case
+tools, and a step budget - run over the same grounded loop the chat assistant
+uses. Every finding must cite a tool result or it is dropped, so no role can
+invent a finding.
+"""
+
+from netforensicai.agents.base import AgentError, AgentFinding, Role, RoleResult, run_role
+from netforensicai.agents.coordinator import MergedFinding, TeamResult, investigate, merge_findings
+from netforensicai.agents.roles import ROLES, all_roles, get_role, resolve_roles
+
+__all__ = [
+    "ROLES",
+    "AgentError",
+    "AgentFinding",
+    "MergedFinding",
+    "Role",
+    "RoleResult",
+    "TeamResult",
+    "all_roles",
+    "get_role",
+    "investigate",
+    "merge_findings",
+    "resolve_roles",
+    "run_role",
+]

netforensicai/agents/base.py

@@ -0,0 +1,226 @@
+"""Foundation for the investigation-team agents.
+
+An agent is one *role*: a mission, a scoped subset of the read-only case tools
+(core/chat.py's CaseTools), and a step budget. It runs over the same grounded
+tool-loop the chat assistant uses, and produces **structured findings** instead
+of free text. The grounding contract is unchanged and non-negotiable: every
+finding must cite a tool result, and a finding that cites something no tool
+returned is dropped - the citation ledger from core/chat.py enforces that
+mechanically, so a role cannot invent a finding.
+
+This module is the reusable substrate. The concrete roles (Network Forensics,
+Host/DFIR, ...) live in roles.py; the coordinator that runs several and merges
+their findings lives in coordinator.py.
+"""
+
+import json
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+from pydantic import BaseModel, ValidationError
+
+from netforensicai.core import ai_assistant
+from netforensicai.core.chat import (
+    TOOL_SPECS,
+    CaseTools,
+    Citation,
+    Ledger,
+    _render,
+    _run_tool,
+)
+
+DEFAULT_ROLE_STEPS = 6
+
+SEVERITIES = ("High", "Medium", "Low", "Info")
+
+
+class AgentError(Exception):
+    """Raised when a role cannot be run at all (bad configuration)."""
+
+
+class AgentFinding(BaseModel):
+    """One structured, evidence-cited finding from a role. Shaped to drop
+    straight onto the investigator-owned Finding model for one-click
+    acceptance."""
+
+    title: str
+    severity: str = "Medium"
+    assessment: str
+    confidence: str = "medium"  # high | medium | low
+    citations: List[Citation] = []
+
+
+@dataclass
+class Role:
+    """A specialist analyst: a name, the mission that scopes its prompt, the
+    tool names it may call (a subset of TOOL_SPECS), and a step budget."""
+
+    name: str
+    slug: str
+    mission: str
+    tools: tuple
+    max_steps: int = DEFAULT_ROLE_STEPS
+
+    def __post_init__(self):
+        unknown = [t for t in self.tools if t not in TOOL_SPECS]
+        if unknown:
+            raise AgentError(f"Role {self.slug!r} lists unknown tools: {unknown}")
+
+    def catalogue(self):
+        return "\n".join(f"- {name}: {TOOL_SPECS[name]}" for name in self.tools)
+
+
+@dataclass
+class RoleResult:
+    """What a role hands back: its cited findings, the tool calls it made, and
+    a note when it found nothing or could not run."""
+
+    role: str
+    slug: str
+    findings: List[AgentFinding] = field(default_factory=list)
+    tool_calls: list = field(default_factory=list)
+    note: Optional[str] = None
+
+    def to_dict(self):
+        return {
+            "role": self.role,
+            "slug": self.slug,
+            "findings": [f.model_dump() for f in self.findings],
+            "tool_calls": self.tool_calls,
+            "note": self.note,
+        }
+
+
+# The protocol the model must speak. Identical grounding rules and citation
+# shape to core/chat.py's SYSTEM_PROMPT - the only difference is the answer
+# form, which here is a list of findings instead of one free-text answer.
+_PROTOCOL = """Reply with a single JSON object and nothing else. Two forms are allowed.
+
+To call a tool:
+{"action": "tool", "tool": "<tool name>", "arguments": {...}}
+
+To report your findings (your final reply):
+{"action": "findings", "findings": [
+  {"title": "<short>", "severity": "High|Medium|Low|Info", "confidence": "high|medium|low",
+   "assessment": "<what the evidence shows and what you infer, kept separate>",
+   "citations": [{"kind": "event|frame|stream|detection", "evidence_id": "EV-0001", "reference": "..."}]}
+]}
+
+Rules that decide whether a finding is usable:
+- Cite ONLY facts a tool returned in this conversation. A finding citing an event, frame, stream or
+  detection that no tool returned is DROPPED. A finding with no citations is dropped.
+- "reference" is the exact identifier from the tool output: an event_id for kind=event, a frame
+  number for kind=frame, a stream index for kind=stream, a rule id for kind=detection.
+- Separate what the evidence shows from what you infer. Phrase inference as possibility
+  ("may indicate", "is consistent with"), never as certainty.
+- Prefer the case's own detections and entities; interpret them, do not invent new ones.
+- If you find nothing in your scope, return {"action": "findings", "findings": []}. That is a
+  correct result, not a failure. Do not guess to fill the list.
+- Retrieve before you report. You have a limited number of tool calls."""
+
+
+def _system_prompt(role):
+    return (
+        f"You are the {role.name} analyst on a digital forensics and incident response (DFIR) "
+        f"investigation team examining ONE case. You cannot see the evidence directly; you retrieve "
+        f"it with tools and may only claim what those tools return.\n\n"
+        f"Your focus: {role.mission}\n\n"
+        f"{_PROTOCOL}\n\n"
+        f"You are not the investigator. Your findings are a starting point for their review."
+    )
+
+
+def _validate_findings(raw_findings, ledger):
+    """Keep only well-formed findings whose citations were all actually
+    retrieved (checked against the ledger, exactly as chat.py checks an
+    answer). Returns (kept_findings, dropped_count)."""
+    kept, dropped = [], 0
+    for item in raw_findings or []:
+        try:
+            finding = AgentFinding.model_validate(item)
+        except ValidationError:
+            dropped += 1
+            continue
+        if not finding.citations:
+            dropped += 1
+            continue
+        if all(ledger.contains(c) for c in finding.citations):
+            kept.append(finding)
+        else:
+            dropped += 1
+    return kept, dropped
+
+
+def run_role(role, case_dir, provider="anthropic", api_key=None, model=None, base_url=None, call=None):
+    """Run one role over `case_dir` and return its cited findings.
+
+    `call(system_prompt, user_prompt) -> dict` overrides the provider call, so
+    the loop can be tested against scripted model behaviour - including a model
+    that cites evidence no tool returned, which must be dropped.
+    """
+    if call is None:
+
+        def call(system_prompt, user_prompt):
+            return ai_assistant.call_model(
+                system_prompt, user_prompt, provider=provider, api_key=api_key, model=model, base_url=base_url
+            )
+
+    ledger = Ledger()
+    tools = CaseTools(case_dir, ledger)
+    system_prompt = _system_prompt(role)
+    transcript = [f"Case directory: {case_dir}", f"Tools available to you:\n{role.catalogue()}"]
+    tool_calls = []
+
+    for step in range(role.max_steps):
+        remaining = role.max_steps - step
+        prompt = "\n\n".join(transcript) + (
+            f"\n\nYou have {remaining} tool call(s) left before you must report findings."
+            if remaining > 1
+            else "\n\nThis is your LAST turn. Report findings now from what you already retrieved."
+        )
+
+        try:
+            raw = call(system_prompt, prompt)
+        except Exception as e:
+            return RoleResult(role.name, role.slug, note=f"provider failed: {e}")
+
+        if not isinstance(raw, dict):
+            transcript.append('Your reply was not a JSON object. Reply with one JSON object.')
+            continue
+
+        action = raw.get("action")
+
+        if action == "tool":
+            name = raw.get("tool")
+            arguments = raw.get("arguments") or {}
+            if name not in role.tools:
+                transcript.append(
+                    f"Tool {name!r} is outside your scope. You may only call: {', '.join(role.tools)}."
+                )
+                continue
+            try:
+                result = _run_tool(tools, name, arguments)
+                rows = len(result) if isinstance(result, list) else 1
+                tool_calls.append({"tool": name, "arguments": arguments, "rows": rows})
+                transcript.append(
+                    f"You called {name}({json.dumps(arguments, default=str)}).\nResult:\n{_render(result)}"
+                )
+            except Exception as e:
+                tool_calls.append({"tool": str(name), "arguments": arguments, "error": str(e)})
+                transcript.append(f"You called {name}({json.dumps(arguments, default=str)}).\nError: {e}")
+            continue
+
+        if action == "findings":
+            kept, dropped = _validate_findings(raw.get("findings"), ledger)
+            note = None
+            if not kept:
+                note = "no findings in scope" if not dropped else "all proposed findings were unciteable and dropped"
+            elif dropped:
+                note = f"{dropped} proposed finding(s) dropped as unciteable"
+            return RoleResult(role.name, role.slug, findings=kept, tool_calls=tool_calls, note=note)
+
+        transcript.append(
+            'Your reply had no valid "action". Reply with {"action":"tool",...} or {"action":"findings",...}.'
+        )
+
+    return RoleResult(role.name, role.slug, tool_calls=tool_calls, note="no findings within the step budget")

netforensicai/agents/coordinator.py

@@ -0,0 +1,150 @@
+"""The Lead Investigator: run the specialist roles over one case and merge
+their cited findings into a single, ranked investigation.
+
+The coordinator does no analysis of its own - that would be an ungrounded
+opinion on top of grounded findings. It only:
+  1. dispatches the requested roles (each already grounded, see base.run_role),
+  2. merges findings that rest on the same evidence (two roles flagging the same
+     event become one, corroborated finding), and
+  3. ranks the result by severity and corroboration.
+
+Every merged finding keeps the union of its citations and the list of roles that
+reported it, so nothing loses its evidence trail.
+"""
+
+from dataclasses import dataclass, field
+from typing import List
+
+from pydantic import BaseModel
+
+from netforensicai.agents.base import RoleResult, run_role
+from netforensicai.agents.roles import all_roles
+
+_SEVERITY_RANK = {"high": 3, "medium": 2, "low": 1, "info": 0}
+
+
+def _sev_rank(severity):
+    return _SEVERITY_RANK.get(str(severity).lower(), 0)
+
+
+class MergedFinding(BaseModel):
+    """A finding after merging: the strongest statement of it, the union of the
+    evidence it rests on, and every role that reported it."""
+
+    title: str
+    severity: str
+    confidence: str
+    assessment: str
+    citations: list  # list[Citation], but kept loose to avoid a re-import cycle
+    reported_by: List[str]
+
+
+@dataclass
+class TeamResult:
+    """The whole team's output: each role's raw result, and the merged, ranked
+    findings across all of them."""
+
+    role_results: List[RoleResult] = field(default_factory=list)
+    findings: List[MergedFinding] = field(default_factory=list)
+
+    def to_dict(self):
+        return {
+            "role_results": [r.to_dict() for r in self.role_results],
+            "findings": [f.model_dump() for f in self.findings],
+        }
+
+
+def _cite_keys(finding):
+    return {(c.kind, c.evidence_id, str(c.reference)) for c in finding.citations}
+
+
+def _merge_group(members):
+    """One merged finding from a group of (role_slug, finding) that share
+    evidence. The lead is the highest-severity member; severity/confidence take
+    the strongest across the group; citations are unioned."""
+    members = sorted(members, key=lambda m: _sev_rank(m[1].severity), reverse=True)
+    _lead_slug, lead = members[0]
+    severity = max((m[1].severity for m in members), key=_sev_rank)
+    confidence = max((m[1].confidence for m in members), key=lambda c: _SEVERITY_RANK.get(str(c).lower(), 0))
+
+    reported_by = []
+    for slug, _ in members:
+        if slug not in reported_by:
+            reported_by.append(slug)
+
+    # Union the citations, de-duplicated on the (kind, evidence_id, reference) key.
+    citations, seen = [], set()
+    for _, finding in members:
+        for c in finding.citations:
+            key = (c.kind, c.evidence_id, str(c.reference))
+            if key not in seen:
+                seen.add(key)
+                citations.append(c)
+
+    if len(members) == 1:
+        assessment = lead.assessment
+    else:
+        # Attribute each distinct assessment to the role(s) that made it.
+        parts, seen_text = [], set()
+        for slug, finding in members:
+            text = finding.assessment.strip()
+            if text and text not in seen_text:
+                seen_text.add(text)
+                parts.append(f"[{slug}] {text}")
+        assessment = "  ".join(parts)
+
+    return MergedFinding(
+        title=lead.title,
+        severity=severity,
+        confidence=confidence,
+        assessment=assessment,
+        citations=citations,
+        reported_by=reported_by,
+    )
+
+
+def merge_findings(role_results):
+    """Group findings that share any citation, merge each group, and rank the
+    result: most severe first, then most-corroborated (reported by more roles),
+    then most evidence."""
+    items = [(r.slug, f) for r in role_results for f in r.findings]
+
+    groups = []  # each: [keys_set, [(slug, finding), ...]]
+    for slug, finding in items:
+        keys = _cite_keys(finding)
+        target = next((g for g in groups if g[0] & keys), None)
+        if target is not None:
+            target[0] |= keys
+            target[1].append((slug, finding))
+        else:
+            groups.append([set(keys), [(slug, finding)]])
+
+    merged = [_merge_group(members) for _keys, members in groups]
+    merged.sort(key=lambda f: (_sev_rank(f.severity), len(f.reported_by), len(f.citations)), reverse=True)
+    return merged
+
+
+def investigate(
+    case_dir,
+    roles=None,
+    provider="anthropic",
+    api_key=None,
+    model=None,
+    base_url=None,
+    call_for=None,
+):
+    """Run `roles` (default: all registered roles) over `case_dir` and return a
+    TeamResult with each role's findings and the merged, ranked set.
+
+    `call_for(role) -> call` overrides the provider call per role, so the team
+    can be tested against scripted model behaviour. In production it is None and
+    each role uses the shared provider settings.
+    """
+    roles = roles if roles is not None else all_roles()
+    role_results = []
+    for role in roles:
+        call = call_for(role) if call_for is not None else None
+        role_results.append(
+            run_role(role, case_dir, provider=provider, api_key=api_key, model=model, base_url=base_url, call=call)
+        )
+    return TeamResult(role_results=role_results, findings=merge_findings(role_results))

netforensicai/agents/roles.py

@@ -0,0 +1,90 @@
+"""The specialist roles of the investigation team.
+
+Each role is a `Role` from base.py: a mission that scopes its system prompt and
+the subset of read-only case tools it may call. Missions tell the model WHAT to
+look for in its domain and to prefer the case's own deterministic detections and
+entities over anything it might infer - the grounding contract (cite or be
+dropped) is enforced by the runner, not by the prose here.
+
+Phase 2 ships the two roles that map to the two evidence domains the tool
+normalizes richly today: network (pcap / Suricata) and host (EVTX / Sysmon).
+The remaining roles (malware/IOC, threat-intel/correlation, reporter) are added
+in a later phase; new roles just register here.
+"""
+
+from netforensicai.agents.base import Role
+
+NETWORK = Role(
+    name="Network Forensics",
+    slug="network",
+    mission=(
+        "Reconstruct what happened on the wire. Concentrate on: connections to "
+        "unusual ports or external hosts; DNS lookups (especially to rare or "
+        "cheap-TLD domains); TLS SNI and HTTP hostnames/URLs; regular-interval "
+        "beaconing that looks like command-and-control; large or one-way "
+        "transfers that look like exfiltration; and any credential seen crossing "
+        "the network in the clear. Start from the case's own detections and "
+        "entities and confirm them against the events, streams and packets; only "
+        "then add what they imply. Do not speculate about host activity you "
+        "cannot see in the network evidence."
+    ),
+    tools=(
+        "list_detections",
+        "list_entities",
+        "search_events",
+        "protocol_summary",
+        "list_streams",
+        "follow_stream",
+        "search_packets",
+    ),
+)
+
+HOST = Role(
+    name="Host & Endpoint (DFIR)",
+    slug="host",
+    mission=(
+        "Reconstruct what happened on the endpoint from Windows Event Log / "
+        "Sysmon events. Concentrate on: process creation and suspicious "
+        "parent-child chains; files written or dropped, especially executables "
+        "and script or key material; persistence; account logons and failed-"
+        "logon bursts; and signs of credential access or lateral movement. Start "
+        "from the case's own detections and host entities (users, hosts, "
+        "processes, files) and confirm them against the events; only then add "
+        "what they imply. Do not speculate about network activity you cannot see "
+        "in the host evidence."
+    ),
+    tools=(
+        "list_detections",
+        "list_entities",
+        "search_events",
+    ),
+)
+
+# Registry, in dispatch order. Later phases append their roles here.
+ROLES = {role.slug: role for role in (NETWORK, HOST)}
+
+
+def all_roles():
+    """Every registered role, in dispatch order."""
+    return list(ROLES.values())
+
+
+def get_role(slug):
+    """Look up a role by slug, or None if there is no such role."""
+    return ROLES.get(slug)
+
+
+def resolve_roles(slugs=None):
+    """The roles named by `slugs` (a list/iterable), or all of them when None.
+
+    Raises KeyError naming the first unknown slug, so a CLI can report it.
+    """
+    if slugs is None:
+        return all_roles()
+    resolved = []
+    for slug in slugs:
+        role = ROLES.get(slug)
+        if role is None:
+            raise KeyError(slug)
+        resolved.append(role)
+    return resolved