n0mad1k
207 lines
8.2 KiB
Python
207 lines
8.2 KiB
Python
"""Interactive Q&A mode for Mosaic — Rich-based chat over extracted intelligence."""
|
|
import logging
|
|
from typing import Optional
|
|
|
|
import anthropic
|
|
from rich.console import Console
|
|
from rich.panel import Panel
|
|
from rich.markdown import Markdown
|
|
from rich.prompt import Prompt
|
|
|
|
from storage.db import MosaicDB
|
|
from .context_builder import ContextBuilder
|
|
|
|
logger = logging.getLogger(__name__)
|
|
console = Console()
|
|
|
|
# Import prompts — handle if not yet created
|
|
try:
|
|
from .prompts import (
|
|
QUERY_SYSTEM_PROMPT, QUERY_PROMPT, EMULATE_PROMPT,
|
|
COMPARE_PROMPT, BRIEF_PROMPT, SCENARIO_PROMPT
|
|
)
|
|
except ImportError:
|
|
QUERY_SYSTEM_PROMPT = "You are an intelligence analyst. Answer questions based on provided documents."
|
|
QUERY_PROMPT = "{context}\n\nQuestion: {question}"
|
|
EMULATE_PROMPT = "Emulate threat actor {actor}:\n{context}"
|
|
COMPARE_PROMPT = "Compare {tool_a} vs {tool_b}:\n{context}"
|
|
BRIEF_PROMPT = "Brief on {topic}:\n{context}"
|
|
SCENARIO_PROMPT = "Scenario for {target}:\n{context}"
|
|
|
|
|
|
class InteractiveAnalyzer:
|
|
"""Interactive Q&A over extracted intelligence corpus."""
|
|
|
|
def __init__(self, db: MosaicDB, api_key: str,
|
|
model: str = "claude-sonnet-4-20250514",
|
|
source: Optional[str] = None):
|
|
self.db = db
|
|
self.client = anthropic.Anthropic(api_key=api_key)
|
|
self.model = model
|
|
self.source = source
|
|
self.context_builder = ContextBuilder(db)
|
|
self.mode = "default"
|
|
self.conversation_history = []
|
|
|
|
def run(self):
|
|
"""Start interactive Q&A session."""
|
|
self._show_welcome()
|
|
|
|
while True:
|
|
try:
|
|
user_input = Prompt.ask("\n[bold green]mosaic[/bold green]")
|
|
except (KeyboardInterrupt, EOFError):
|
|
console.print("\n[dim]Goodbye.[/dim]")
|
|
break
|
|
|
|
user_input = user_input.strip()
|
|
if not user_input:
|
|
continue
|
|
|
|
# Handle commands
|
|
if user_input.lower() in ('exit', 'quit', 'q'):
|
|
console.print("[dim]Goodbye.[/dim]")
|
|
break
|
|
elif user_input.lower() == 'help':
|
|
self._show_help()
|
|
continue
|
|
elif user_input.lower().startswith('mode '):
|
|
self.mode = user_input[5:].strip()
|
|
console.print(f"[cyan]Mode set to: {self.mode}[/cyan]")
|
|
continue
|
|
elif user_input.lower().startswith('source '):
|
|
self.source = user_input[7:].strip()
|
|
console.print(f"[cyan]Source filter: {self.source}[/cyan]")
|
|
continue
|
|
elif user_input.lower() == 'clear':
|
|
self.conversation_history = []
|
|
console.print("[dim]Conversation cleared.[/dim]")
|
|
continue
|
|
elif user_input.lower() == 'status':
|
|
self._show_status()
|
|
continue
|
|
|
|
# Process query based on mode
|
|
self._handle_query(user_input)
|
|
|
|
def _handle_query(self, query: str):
|
|
"""Process a query in the current mode."""
|
|
with console.status("[bold cyan]Searching corpus...[/bold cyan]"):
|
|
# Build context based on mode
|
|
if self.mode == 'emulate':
|
|
context = self.context_builder.build_context(query, self.source)
|
|
prompt = EMULATE_PROMPT.format(actor=query, context=context['context_text'])
|
|
elif self.mode == 'compare' and ' vs ' in query:
|
|
parts = query.split(' vs ', 1)
|
|
tool_a = parts[0].strip()
|
|
tool_b = parts[1].strip()
|
|
ctx_a = self.context_builder.build_tool_context(tool_a)
|
|
ctx_b = self.context_builder.build_tool_context(tool_b)
|
|
context = {'context_text': ctx_a + '\n\n' + ctx_b, 'documents': [], 'total_tokens': 0}
|
|
prompt = COMPARE_PROMPT.format(tool_a=tool_a, tool_b=tool_b, context=context['context_text'])
|
|
elif self.mode == 'brief':
|
|
context = self.context_builder.build_context(query, self.source)
|
|
prompt = BRIEF_PROMPT.format(topic=query, context=context['context_text'])
|
|
elif self.mode == 'scenario':
|
|
context = self.context_builder.build_context(query, self.source)
|
|
prompt = SCENARIO_PROMPT.format(target=query, context=context['context_text'])
|
|
else:
|
|
context = self.context_builder.build_context(query, self.source)
|
|
prompt = QUERY_PROMPT.format(context=context['context_text'], question=query)
|
|
|
|
if not context.get('context_text'):
|
|
console.print("[yellow]No relevant documents found for your query.[/yellow]")
|
|
return
|
|
|
|
# Show source count
|
|
doc_count = len(context.get('documents', []))
|
|
console.print(f"[dim]Found {doc_count} relevant documents, "
|
|
f"~{context.get('total_tokens', 0):,} tokens of context[/dim]")
|
|
|
|
# Call Claude
|
|
with console.status("[bold cyan]Analyzing...[/bold cyan]"):
|
|
try:
|
|
messages = list(self.conversation_history)
|
|
messages.append({"role": "user", "content": prompt})
|
|
|
|
response = self.client.messages.create(
|
|
model=self.model,
|
|
max_tokens=4096,
|
|
system=QUERY_SYSTEM_PROMPT,
|
|
messages=messages,
|
|
)
|
|
|
|
answer = response.content[0].text
|
|
|
|
# Save to conversation history
|
|
self.conversation_history.append({"role": "user", "content": prompt})
|
|
self.conversation_history.append({"role": "assistant", "content": answer})
|
|
|
|
# Trim history if too long
|
|
if len(self.conversation_history) > 20:
|
|
self.conversation_history = self.conversation_history[-10:]
|
|
|
|
except anthropic.APIError as e:
|
|
console.print(f"[red]API error: {e}[/red]")
|
|
return
|
|
|
|
# Display response
|
|
console.print()
|
|
console.print(Panel(
|
|
Markdown(answer),
|
|
title="[bold]Analysis[/bold]",
|
|
border_style="cyan",
|
|
))
|
|
|
|
# Show citations
|
|
if context.get('documents'):
|
|
console.print("\n[dim]Sources:[/dim]")
|
|
for doc in context['documents'][:5]:
|
|
console.print(f" [dim]- {doc['source']}: {doc['title']}", markup=False)
|
|
|
|
def _show_welcome(self):
|
|
"""Show welcome message."""
|
|
welcome = """[bold cyan]MOSAIC Interactive Intelligence Query[/bold cyan]
|
|
|
|
Ask questions about the extracted intelligence corpus.
|
|
Type [bold]help[/bold] for commands, [bold]exit[/bold] to quit."""
|
|
|
|
if self.source:
|
|
welcome += f"\n[dim]Source filter: {self.source}[/dim]"
|
|
|
|
console.print(Panel(welcome, border_style="cyan"))
|
|
|
|
def _show_help(self):
|
|
"""Show help text."""
|
|
help_text = """[bold]Commands:[/bold]
|
|
help Show this help
|
|
exit/quit Leave interactive mode
|
|
clear Clear conversation history
|
|
status Show corpus statistics
|
|
mode <mode> Set query mode:
|
|
default - General Q&A
|
|
emulate - Threat actor emulation playbook
|
|
compare - Compare tools (use 'tool_a vs tool_b')
|
|
brief - Intelligence briefing
|
|
scenario - Adversary simulation scenario
|
|
source <name> Filter to specific source (or 'all')
|
|
|
|
[bold]Tips:[/bold]
|
|
- Ask specific questions for better results
|
|
- Use source names (vault7, cablegate) to focus queries
|
|
- Emulate mode: type an actor name to get their TTP playbook"""
|
|
console.print(Panel(help_text, border_style="cyan"))
|
|
|
|
def _show_status(self):
|
|
"""Show corpus statistics."""
|
|
stats = self.db.get_status_counts(self.source)
|
|
tokens = self.db.get_total_tokens(self.source)
|
|
|
|
console.print(f"\n[bold]Corpus Statistics[/bold]")
|
|
console.print(f" Documents: {stats.get('total', 0)}")
|
|
console.print(f" Parsed: {stats.get('parsing', {}).get('PARSED', 0)}")
|
|
console.print(f" Tokens used: {tokens.get('total_tokens', 0):,}")
|
|
|
|
if self.source:
|
|
console.print(f" [dim]Filtered to: {self.source}[/dim]")
|