Update app.py

app.py

@@ -3,9 +3,9 @@
 Parrot OSINT MCP – Gradio Frontend
 
 Modes:
-- "OSINT Dashboard" (multi-tool, opinionated)
-- "MCP Bridge" (raw tool_name + JSON args → JSON result)
-- "Analyst Copilot" (streaming LLM with OSINT context injection)
+- OSINT Dashboard (deterministic intelligence)
+- MCP Bridge (raw tool access)
+- Analyst Copilot (LLM interpretive intelligence)
 """
 
 import json
@@ -16,263 +16,238 @@ import gradio as gr
 from huggingface_hub import InferenceClient
 
 # ---------------------------------------------------------------------
-# Task registry: adapt this to your actual task API
+# Task Registry (auto-loads your MCP tasks)
 # ---------------------------------------------------------------------
 
 TASK_REGISTRY: Dict[str, Any] = {}
 
-def _register_tasks() -> None:
-    def _try_register(name: str, module_name: str):
+def _register_tasks():
+    def _try(name, module):
         try:
-            module = __import__(f"tasks.{module_name}", fromlist=["*"])
-            fn = getattr(module, "run", None)
+            m = __import__(f"tasks.{module}", fromlist=["*"])
+            fn = getattr(m, "run", None)
             if callable(fn):
                 TASK_REGISTRY[name] = fn
         except Exception:
             pass
 
-    _try_register("lookup_ip", "lookup_ip")
-    _try_register("lookup_domain", "lookup_domain")
-    _try_register("lookup_hash", "lookup_hash")
-    _try_register("correlate_iocs", "correlate_iocs")
-    _try_register("generate_report", "generate_report")
-    _try_register("enrich_entity", "enrich_entity")
-    _try_register("mitre_map", "mitre_map")
-    _try_register("quickscan", "quickscan")
+    _try("lookup_ip", "lookup_ip")
+    _try("lookup_domain", "lookup_domain")
+    _try("lookup_hash", "lookup_hash")
+    _try("correlate_iocs", "correlate_iocs")
+    _try("generate_report", "generate_report")
+    _try("enrich_entity", "enrich_entity")
+    _try("mitre_map", "mitre_map")
+    _try("quickscan", "quickscan")
 
 _register_tasks()
 
-
 # ---------------------------------------------------------------------
-# Core execution helpers
+# Core Task Execution
 # ---------------------------------------------------------------------
 
-def call_task(tool_name: str, payload: Dict[str, Any]) -> Dict[str, Any]:
-    fn = TASK_REGISTRY.get(tool_name)
+def call_task(name: str, payload: Dict[str, Any]):
+    fn = TASK_REGISTRY.get(name)
     if not fn:
-        return {
-            "error": f"Unknown tool '{tool_name}'. Registered tools: {sorted(TASK_REGISTRY.keys())}"
-        }
+        return {"error": f"Unknown tool '{name}'."}
 
     try:
-        result = fn(**payload)
-        if not isinstance(result, dict):
-            result = {"result": result}
-        return result
-    except Exception as exc:
-        return {
-            "error": f"Exception in tool '{tool_name}': {exc}",
-            "traceback": traceback.format_exc(),
-        }
-
-
-def format_result_for_ui(result: Dict[str, Any]) -> Dict[str, str]:
-    pretty_json = json.dumps(result, indent=2, default=str)
-
-    markdown = result.get("markdown") or result.get("report") or ""
-    if not markdown and "summary" in result:
-        markdown = f"## Summary\n\n{result['summary']}"
-
-    mitre = json.dumps(result.get("mitre", ""), indent=2, default=str) if result.get("mitre") else ""
-    stix  = json.dumps(result.get("stix", ""), indent=2, default=str) if result.get("stix")  else ""
-    sarif = json.dumps(result.get("sarif", ""), indent=2, default=str) if result.get("sarif") else ""
+        res = fn(**payload)
+        if not isinstance(res, dict):
+            res = {"result": res}
+        return res
+    except Exception as e:
+        return {"error": str(e), "traceback": traceback.format_exc()}
+
+
+def normalize_result(res: Dict[str, Any]):
+    """Formats UI fields cleanly."""
+    pretty = json.dumps(res, indent=2, default=str)
+    summary = res.get("summary", "")
+    markdown = res.get("markdown") or res.get("report") or ""
+    if not markdown and summary:
+        markdown = f"## Summary\n\n{summary}"
 
     return {
-        "summary": result.get("summary", ""),
+        "summary": summary,
         "markdown": markdown,
-        "json": pretty_json,
-        "mitre": mitre,
-        "stix": stix,
-        "sarif": sarif,
+        "json": pretty,
+        "mitre": json.dumps(res.get("mitre", ""), indent=2, default=str) if res.get("mitre") else "",
+        "stix": json.dumps(res.get("stix", ""), indent=2, default=str) if res.get("stix") else "",
+        "sarif": json.dumps(res.get("sarif", ""), indent=2, default=str) if res.get("sarif") else "",
     }
 
-
 # ---------------------------------------------------------------------
-# MODE C — ANALYST COPILOT (LLM)
+# ANALYST COPILOT (LLM)
 # ---------------------------------------------------------------------
 
-def respond(message, history, system_message, model, hf_token, temperature, top_p, max_tokens):
-    """
-    Streaming LLM response using HuggingFace InferenceClient.
-    Supports injecting OSINT task results into the conversation.
-    """
-    client = InferenceClient(
-        token=hf_token.token,
-        model=model,
-    )
-
-    messages = [{"role": "system", "content": system_message}]
-    messages.extend(history)
-    messages.append({"role": "user", "content": message})
-
-    response_text = ""
-
+def respond(
+    message,
+    history,
+    system_prompt,
+    model_name,
+    hf_token,
+    temperature,
+    top_p,
+    max_tokens,
+):
+    """Streaming response from WhiteRabbit Neo or Cybertron."""
+    client = InferenceClient(model=model_name, token=hf_token.token)
+
+    msgs = [{"role": "system", "content": system_prompt}]
+    msgs.extend(history)
+    msgs.append({"role": "user", "content": message})
+
+    buf = ""
     for chunk in client.chat_completion(
-        messages=messages,
+        messages=msgs,
+        max_tokens=max_tokens,
         temperature=temperature,
         top_p=top_p,
-        max_tokens=max_tokens,
-        stream=True
+        stream=True,
     ):
         delta = chunk.choices[0].delta.content
         if delta:
-            response_text += delta
-            yield response_text
-
-
-def inject_osint_context(history, task_result: Dict[str, Any]):
-    """
-    Inject JSON + summary + MITRE mappings directly into the chat history.
-    """
-    pretty = json.dumps(task_result, indent=2, default=str)
-    blob = f"""
-### OSINT Result Injected:
-
-{pretty}
-
-"""
-
-    history.append({"role": "system", "content": blob})
+            buf += delta
+            yield buf
+
+def inject_osint(history, osint_obj):
+    """Inject raw JSON results into the chat context."""
+    pretty = json.dumps(osint_obj, indent=2, default=str)
+    history.append({
+        "role": "system",
+        "content": f"### Injected OSINT Result\n```\n{pretty}\n```"
+    })
     return history
 
-
 # ---------------------------------------------------------------------
-# Dashboard callbacks (Mode B)
+# OSINT Dashboard Callbacks
 # ---------------------------------------------------------------------
 
 def ui_lookup_ip(ip, enrich, mitre):
     raw = call_task("lookup_ip", {"ip": ip, "enrich": enrich, "map_mitre": mitre})
-    normal = format_result_for_ui(raw)
-    return normal["summary"], normal["markdown"], normal["json"], normal["mitre"], normal["stix"], raw
-
+    norm = normalize_result(raw)
+    return norm["summary"], norm["markdown"], norm["json"], norm["mitre"], norm["stix"], raw
 
 def ui_lookup_domain(domain, enrich, mitre):
     raw = call_task("lookup_domain", {"domain": domain, "enrich": enrich, "map_mitre": mitre})
-    normal = format_result_for_ui(raw)
-    return normal["summary"], normal["markdown"], normal["json"], normal["mitre"], normal["stix"], raw
-
+    norm = normalize_result(raw)
+    return norm["summary"], norm["markdown"], norm["json"], norm["mitre"], norm["stix"], raw
 
 def ui_lookup_hash(h, ht, enrich, mitre):
     raw = call_task("lookup_hash", {"hash": h, "hash_type": ht, "enrich": enrich, "map_mitre": mitre})
-    normal = format_result_for_ui(raw)
-    return normal["summary"], normal["markdown"], normal["json"], normal["mitre"], normal["stix"], raw
-
+    norm = normalize_result(raw)
+    return norm["summary"], norm["markdown"], norm["json"], norm["mitre"], norm["stix"], raw
 
 def ui_correlate_iocs(iocs):
-    parsed = [l.strip() for l in iocs.splitlines() if l.strip()]
-    raw = call_task("correlate_iocs", {"iocs": parsed})
-    normal = format_result_for_ui(raw)
-    return normal["summary"], normal["markdown"], normal["json"], normal["mitre"], raw
-
+    lst = [x.strip() for x in iocs.splitlines() if x.strip()]
+    raw = call_task("correlate_iocs", {"iocs": lst})
+    norm = normalize_result(raw)
+    return norm["summary"], norm["markdown"], norm["json"], norm["mitre"], raw
 
 def ui_quickscan(target):
     raw = call_task("quickscan", {"target": target})
-    normal = format_result_for_ui(raw)
-    return normal["summary"], normal["markdown"], normal["json"], raw
-
+    norm = normalize_result(raw)
+    return norm["summary"], norm["markdown"], norm["json"], raw
 
 # ---------------------------------------------------------------------
-# MCP Bridge (Mode D)
+# MCP Bridge
 # ---------------------------------------------------------------------
 
-def ui_mcp_bridge(tool, args_json):
+def ui_bridge(tool, args_json):
     try:
         payload = json.loads(args_json)
-    except Exception as exc:
-        err = {"error": f"Invalid JSON: {exc}"}
-        return json.dumps(err, indent=2), "", err
-
+    except Exception as e:
+        return json.dumps({"error": str(e)}, indent=2), "", {}
     raw = call_task(tool, payload)
-    normal = format_result_for_ui(raw)
-    return normal["json"], normal["markdown"], raw
-
+    norm = normalize_result(raw)
+    return norm["json"], norm["markdown"], raw
 
 # ---------------------------------------------------------------------
-# UI — Now with Analyst Copilot
+# BUILD UI
 # ---------------------------------------------------------------------
 
-def build_interface() -> gr.Blocks:
+def build_interface():
     with gr.Blocks(title="Parrot OSINT MCP Console") as demo:
         gr.Markdown("# Parrot OSINT MCP Console")
 
-        # Store OSINT task results for injection into the Copilot
-        osint_result_state = gr.State([])
+        osint_state = gr.State({})
 
-        # ------------------------------------------
-        # MODE B — Dashboard
-        # ------------------------------------------
+        # -------------------------
+        # OSINT Dashboard
+        # -------------------------
         with gr.Tab("OSINT Dashboard"):
-            with gr.Tab("IP Lookup"):
+            with gr.Tab("IP"):
                 ip = gr.Textbox(label="IP Address")
-                enrich = gr.Checkbox(value=True, label="Enrichment")
-                mitre = gr.Checkbox(value=True, label="MITRE mapping")
-                btn = gr.Button("Run")
-                summary = gr.Textbox(label="Summary")
+                enrich = gr.Checkbox(value=True, label="Enrich")
+                mitre = gr.Checkbox(value=True, label="MITRE Map")
+                run = gr.Button("Run IP Lookup")
+
+                s = gr.Textbox(label="Summary")
                 md = gr.Markdown()
                 js = gr.Code(language="json")
                 mt = gr.Code(language="json")
                 st = gr.Code(language="json")
 
-                btn.click(
-                    ui_lookup_ip,
-                    inputs=[ip, enrich, mitre],
-                    outputs=[summary, md, js, mt, st, osint_result_state],
-                )
+                run.click(ui_lookup_ip, [ip, enrich, mitre], [s, md, js, mt, st, osint_state])
 
-            # You already know: similar tabs for domain, hash, correlation, quickscan
-            # (keeping focus on Copilot integration)
+            # Add other tabs (Domain, Hash, etc.)
+            # Your earlier implementation plugs in cleanly.
 
-        # ------------------------------------------
-        # MODE D — MCP Bridge
-        # ------------------------------------------
+        # -------------------------
+        # MCP Bridge
+        # -------------------------
         with gr.Tab("MCP Bridge"):
             tool = gr.Dropdown(sorted(TASK_REGISTRY.keys()))
-            args = gr.Code(label="Args JSON")
+            args = gr.Code(language="json")
+            btn = gr.Button("Call Tool")
             out_js = gr.Code(language="json")
             out_md = gr.Markdown()
 
-            bridge_btn = gr.Button("Call Tool")
-            bridge_btn.click(
-                ui_mcp_bridge,
-                inputs=[tool, args],
-                outputs=[out_js, out_md, osint_result_state],
-            )
+            btn.click(ui_bridge, [tool, args], [out_js, out_md, osint_state])
 
-        # ------------------------------------------
-        # MODE C — Analyst Copilot
-        # ------------------------------------------
+        # -------------------------
+        # Analyst Copilot
+        # -------------------------
         with gr.Tab("Analyst Copilot"):
-            gr.Markdown("### Streaming TI Assistant with OSINT Context Injection")
+            gr.Markdown("### WhiteRabbit Neo + Cybertron TI Assistant")
 
-            system_msg = gr.Textbox(
+            system_prompt = gr.Textbox(
                 label="System Prompt",
-                value=("You are a threat intelligence analyst. "
-                       "You think slowly, explain clearly, identify TTPs, "
-                       "and recommend next investigative steps."),
+                value=(
+                    "You are a threat intelligence analyst. "
+                    "You classify TTPs, map MITRE ATT&CK, and provide investigation guidance."
+                ),
             )
 
-            model = gr.Textbox(
-                label="HF Model (e.g., openai/gpt-oss-20b)",
-                value="openai/gpt-oss-20b",
+            model_select = gr.Dropdown(
+                label="LLM Model",
+                choices=[
+                    "berkeley-nest/WhiteRabbitNeo-8B",
+                    "cybertronai/cybertron-1.1-1b",
+                    "cybertronai/cybertron-1.1-7b",
+                    "cybertronai/cybertron-1.1-32b"
+                ],
+                value="berkeley-nest/WhiteRabbitNeo-8B",
             )
 
             chatbot = gr.ChatInterface(
                 respond,
+                type="messages",
                 additional_inputs=[
-                    system_msg,
-                    model,
+                    system_prompt,
+                    model_select,
                     gr.OAuthToken(label="HF Token"),
                     gr.Slider(0.1, 2.0, value=0.7, step=0.1, label="Temperature"),
                     gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p"),
-                    gr.Slider(1, 2048, value=512, step=1, label="Max Tokens"),
+                    gr.Slider(32, 4096, value=512, step=32, label="Max Tokens"),
                 ],
-                type="messages",
             )
 
-            inject_btn = gr.Button("Inject Latest OSINT Result")
+            inject_btn = gr.Button("Inject Last OSINT Result into Chat")
             inject_btn.click(
-                inject_osint_context,
-                inputs=[chatbot._chatbot_state, osint_result_state],
+                inject_osint,
+                inputs=[chatbot._chatbot_state, osint_state],
                 outputs=[chatbot._chatbot_state],
             )