diff --git a/analysis/__init__.py b/analysis/__init__.py index e69de29..335e3ce 100644 --- a/analysis/__init__.py +++ b/analysis/__init__.py @@ -0,0 +1,3 @@ +from .claude_batch import BatchAnalyzer +from .claude_interactive import InteractiveAnalyzer +from .context_builder import ContextBuilder diff --git a/analysis/claude_batch.py b/analysis/claude_batch.py new file mode 100644 index 0000000..d602931 --- /dev/null +++ b/analysis/claude_batch.py @@ -0,0 +1,258 @@ +"""Batch document analysis pipeline for Mosaic.""" +import logging +from typing import Optional +from datetime import datetime + +from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn +from rich.console import Console + +from storage.db import MosaicDB +from storage.models import Document +from storage.jsonl_store import JSONLStore +from utils.api_limiter import APIRateLimiter, CircuitBreakerOpen +from utils.text_utils import estimate_tokens + +logger = logging.getLogger(__name__) +console = Console() + +EXTRACTOR_DOMAINS = { + 'tools': 'tool', + 'ttps': 'ttp', + 'tradecraft': 'tradecraft', + 'surveillance': 'surveillance', + 'infrastructure': 'infrastructure', + 'entities': 'entity', +} + + +class BatchAnalyzer: + """Run batch extraction pipeline across documents.""" + + def __init__(self, db: MosaicDB, api_key: str, + model: str = "claude-sonnet-4-20250514", + rate_limiter: Optional[APIRateLimiter] = None): + self.db = db + self.api_key = api_key + self.model = model + self.rate_limiter = rate_limiter or APIRateLimiter() + self._extractors = {} + + def _get_extractor(self, domain: str): + """Lazy-load extractor by domain name.""" + if domain not in self._extractors: + from extractors.ttp_extractor import TTPExtractor + from extractors.tool_extractor import ToolExtractor + from extractors.tradecraft import TradecraftExtractor + from extractors.surveillance import SurveillanceExtractor + from extractors.infrastructure import InfrastructureExtractor + from extractors.entity_extractor import EntityExtractor + + extractor_classes = { + 'ttp': TTPExtractor, + 'tool': ToolExtractor, + 'tradecraft': TradecraftExtractor, + 'surveillance': SurveillanceExtractor, + 'infrastructure': InfrastructureExtractor, + 'entity': EntityExtractor, + } + + cls = extractor_classes.get(domain) + if cls: + if domain == 'infrastructure': + self._extractors[domain] = cls( + db=self.db, api_key=self.api_key, + model=self.model, rate_limiter=self.rate_limiter, + collection="", doc_title="", doc_date=None, + ) + else: + self._extractors[domain] = cls( + db=self.db, api_key=self.api_key, + model=self.model, rate_limiter=self.rate_limiter, + ) + return self._extractors.get(domain) + + def estimate(self, source: Optional[str] = None, + domains: Optional[list[str]] = None, + since: Optional[str] = None, + until: Optional[str] = None, + since_last_run: bool = False) -> dict: + """Estimate token costs without running extraction (dry-run). + + Returns dict with document count, chunk count, estimated tokens. + """ + domains = domains or list(EXTRACTOR_DOMAINS.values()) + docs = self._get_target_documents(source, since, until, since_last_run) + + total_chunks = 0 + total_input_tokens = 0 + + for doc in docs: + for domain in domains: + extractor = self._get_extractor(domain) + if extractor: + est = extractor.estimate_tokens(doc.text, doc.doc_type) + total_chunks += est['chunks'] + total_input_tokens += est['estimated_input_tokens'] + + return { + 'document_count': len(docs), + 'domain_count': len(domains), + 'total_chunks': total_chunks, + 'estimated_input_tokens': total_input_tokens, + 'estimated_output_tokens': total_chunks * 2000, + 'estimated_total_tokens': total_input_tokens + (total_chunks * 2000), + } + + def run(self, source: Optional[str] = None, + domains: Optional[list[str]] = None, + since: Optional[str] = None, + until: Optional[str] = None, + since_last_run: bool = False, + retry_failed: bool = False) -> dict: + """Run batch extraction pipeline. + + Returns dict with counts of extracted items per domain. + """ + domains = domains or list(EXTRACTOR_DOMAINS.values()) + docs = self._get_target_documents(source, since, until, since_last_run) + + if not docs: + console.print("[yellow]No documents to process[/yellow]") + return {} + + results = {d: 0 for d in domains} + + with Progress( + SpinnerColumn(), + TextColumn("[progress.description]{task.description}"), + BarColumn(), + TaskProgressColumn(), + ) as progress: + for domain in domains: + extractor = self._get_extractor(domain) + if not extractor: + console.print(f"[yellow]Unknown domain: {domain}[/yellow]") + continue + + task = progress.add_task(f"Extracting {domain}", total=len(docs)) + + for doc in docs: + # Check if already done + status = self.db.get_extraction_status(doc.id, extractor.EXTRACTOR_NAME) + if status and status.status == 'DONE' and not retry_failed: + progress.advance(task) + continue + if status and status.status == 'FAILED' and not retry_failed: + progress.advance(task) + continue + + try: + # Update infrastructure extractor context + if domain == 'infrastructure' and hasattr(extractor, 'collection'): + extractor.collection = doc.source + extractor.doc_title = doc.title + extractor.doc_date = doc.date + + count = extractor.extract_from_document( + doc_id=doc.id, + text=doc.text, + doc_type=doc.doc_type, + source=doc.source, + metadata=doc.metadata, + ) + results[domain] += count + + except CircuitBreakerOpen: + console.print(f"[red]Circuit breaker open — stopping {domain} extraction[/red]") + break + except Exception as e: + logger.error("Extraction error for %s on doc %s: %s", + domain, doc.id[:12], e) + + progress.advance(task) + + return results + + def run_full_pipeline(self, source: str, limit: Optional[int] = None) -> dict: + """Run full pipeline: collect → parse → extract → report. + + Returns summary dict. + """ + from collectors import load_profile, CustomSource + from parsers import ParserManager + from utils.cache import DownloadCache + from utils.http_limiter import HTTPRateLimiter + + summary = {'collect': {}, 'parse': {}, 'extract': {}, 'export': []} + + # Step 1: Collect + console.print(f"\n[bold cyan]Step 1/4: Collecting from {source}[/bold cyan]") + try: + profile = load_profile(source) + cache = DownloadCache() + collector = CustomSource(profile, self.db, cache) + summary['collect'] = collector.collect(limit=limit) + except Exception as e: + console.print(f"[red]Collection error: {e}[/red]") + return summary + + # Step 2: Parse + console.print(f"\n[bold cyan]Step 2/4: Parsing documents[/bold cyan]") + parser = ParserManager() + docs = self.db.get_documents_by_source(source, status='CACHED') + parsed = 0 + failed = 0 + + with Progress(SpinnerColumn(), TextColumn("{task.description}"), + BarColumn(), TaskProgressColumn()) as progress: + task = progress.add_task("Parsing", total=len(docs)) + for doc in docs: + if doc.parse_status in ('PARSED', 'NEEDS_OCR'): + progress.advance(task) + continue + + self.db.update_document_status(doc.id, parse_status='IN_PROGRESS') + result = parser.parse(doc.raw_path, doc.doc_type, doc.content_hash) + + if result.success: + doc.text = result.text + doc.title = result.title or doc.title + doc.char_count = result.char_count + doc.metadata.update(result.metadata) + status = 'NEEDS_OCR' if result.needs_ocr else 'PARSED' + doc.parse_status = status + self.db.insert_document(doc) + self.db.update_document_status(doc.id, parse_status=status) + parsed += 1 + else: + self.db.update_document_status(doc.id, parse_status='PARSE_FAILED') + failed += 1 + + progress.advance(task) + + summary['parse'] = {'parsed': parsed, 'failed': failed} + + # Step 3: Extract + console.print(f"\n[bold cyan]Step 3/4: Running extraction[/bold cyan]") + summary['extract'] = self.run(source=source) + + # Step 4: Export + console.print(f"\n[bold cyan]Step 4/4: Exporting results[/bold cyan]") + store = JSONLStore(self.db) + summary['export'] = store.export() + + return summary + + def _get_target_documents(self, source: Optional[str] = None, + since: Optional[str] = None, + until: Optional[str] = None, + since_last_run: bool = False) -> list[Document]: + """Get documents to process based on filters.""" + if since or until: + return self.db.get_documents_by_date_range(source=source, since=since, until=until) + elif since_last_run: + return self.db.get_unextracted_documents(source=source, since_last_run=True) + elif source: + return self.db.get_documents_by_source(source, parse_status='PARSED') + else: + return self.db.get_unextracted_documents() diff --git a/analysis/claude_interactive.py b/analysis/claude_interactive.py new file mode 100644 index 0000000..5758ab8 --- /dev/null +++ b/analysis/claude_interactive.py @@ -0,0 +1,206 @@ +"""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 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 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]") diff --git a/analysis/context_builder.py b/analysis/context_builder.py new file mode 100644 index 0000000..19b386b --- /dev/null +++ b/analysis/context_builder.py @@ -0,0 +1,168 @@ +"""Context builder for assembling relevant document chunks for Claude Q&A.""" +import logging +from typing import Optional + +from storage.db import MosaicDB +from utils.text_utils import estimate_tokens + +logger = logging.getLogger(__name__) + +DEFAULT_TOKEN_BUDGET = 30000 +DEFAULT_MAX_RESULTS = 20 + + +class ContextBuilder: + """Build context from FTS5 search results for Claude queries. + + Uses BM25 ranking from FTS5, assembles chunks within a token budget, + and tracks source citations. + """ + + def __init__(self, db: MosaicDB, token_budget: int = DEFAULT_TOKEN_BUDGET, + max_results: int = DEFAULT_MAX_RESULTS): + self.db = db + self.token_budget = token_budget + self.max_results = max_results + + def build_context(self, query: str, source: Optional[str] = None) -> dict: + """Search documents and assemble context within token budget. + + Args: + query: Search query string + source: Optional source filter + + Returns: + dict with: + 'context_text': assembled context string with citations + 'documents': list of matched document metadata + 'total_tokens': estimated token count + 'truncated': whether results were truncated + """ + # Search with FTS5 + docs = self.db.search_documents(query, source=source, limit=self.max_results) + + if not docs: + return { + 'context_text': '', + 'documents': [], + 'total_tokens': 0, + 'truncated': False, + } + + # Assemble context within budget + context_parts = [] + used_tokens = 0 + included_docs = [] + truncated = False + + for doc in docs: + # Build document excerpt + excerpt = self._format_document_excerpt(doc) + excerpt_tokens = estimate_tokens(excerpt) + + if used_tokens + excerpt_tokens > self.token_budget: + # Try to fit a truncated version + remaining_tokens = self.token_budget - used_tokens + if remaining_tokens > 500: # Worth including a partial + truncated_excerpt = self._truncate_to_tokens(excerpt, remaining_tokens) + context_parts.append(truncated_excerpt) + used_tokens += remaining_tokens + included_docs.append(self._doc_metadata(doc)) + truncated = True + break + + context_parts.append(excerpt) + used_tokens += excerpt_tokens + included_docs.append(self._doc_metadata(doc)) + + context_text = '\n\n---\n\n'.join(context_parts) + + return { + 'context_text': context_text, + 'documents': included_docs, + 'total_tokens': used_tokens, + 'truncated': truncated, + } + + def build_entity_context(self, entity_name: str) -> str: + """Build context focused on a specific entity across all sources.""" + docs = self.db.search_documents(entity_name, limit=self.max_results) + + parts = [] + used_tokens = 0 + for doc in docs: + excerpt = self._format_document_excerpt(doc) + tokens = estimate_tokens(excerpt) + if used_tokens + tokens > self.token_budget: + break + parts.append(excerpt) + used_tokens += tokens + + return '\n\n---\n\n'.join(parts) + + def build_tool_context(self, tool_name: str) -> str: + """Build context focused on a specific tool/implant.""" + # Search both FTS and extracted tools + docs = self.db.search_documents(tool_name, limit=self.max_results) + + parts = [] + used_tokens = 0 + + # Add extracted tool data if available + tools = self.db.get_all_tools() + for tool in tools: + if tool_name.lower() in tool.name.lower() or \ + any(tool_name.lower() in a.lower() for a in tool.aliases): + tool_text = f"[Extracted Tool: {tool.name}]\n" + tool_text += f"Capability: {tool.capability}\n" + tool_text += f"Description: {tool.description}\n" + tool_text += f"Platforms: {', '.join(tool.target_platforms)}\n" + tool_text += f"CVEs: {', '.join(tool.cves)}\n" + tool_text += f"Source context: {tool.source_context}\n" + parts.append(tool_text) + used_tokens += estimate_tokens(tool_text) + + # Add document excerpts + for doc in docs: + excerpt = self._format_document_excerpt(doc) + tokens = estimate_tokens(excerpt) + if used_tokens + tokens > self.token_budget: + break + parts.append(excerpt) + used_tokens += tokens + + return '\n\n---\n\n'.join(parts) + + def _format_document_excerpt(self, doc) -> str: + """Format a document for context inclusion with citation info.""" + header = f"[Source: {doc.source} | Title: {doc.title}" + if doc.date: + header += f" | Date: {doc.date}" + if doc.classification: + header += f" | Classification: {doc.classification}" + header += f" | ID: {doc.id[:12]}]" + + # Truncate text if needed + text = doc.text + max_chars = (self.token_budget // len(self.db.search_documents("", limit=1) or [doc])) * 4 + if len(text) > max_chars: + text = text[:max_chars] + "\n[...truncated...]" + + return f"{header}\n\n{text}" + + def _truncate_to_tokens(self, text: str, max_tokens: int) -> str: + """Truncate text to approximately max_tokens.""" + max_chars = max_tokens * 4 # Rough chars-per-token estimate + if len(text) <= max_chars: + return text + return text[:max_chars] + "\n[...truncated...]" + + def _doc_metadata(self, doc) -> dict: + """Extract metadata dict from a document for citation tracking.""" + return { + 'id': doc.id, + 'title': doc.title, + 'source': doc.source, + 'date': doc.date, + 'doc_type': doc.doc_type, + }