diff --git a/analysis/prompts.py b/analysis/prompts.py new file mode 100644 index 0000000..72390c3 --- /dev/null +++ b/analysis/prompts.py @@ -0,0 +1,224 @@ +"""Prompt templates for Claude-powered extraction in Mosaic. + +All document content is wrapped in XML tags to prevent +prompt injection. Claude responses must be valid JSON matching +the specified schema. +""" + +# --- TTP Extraction --- + +TTP_SYSTEM_PROMPT = """You are an intelligence analyst specializing in cyber threat analysis and MITRE ATT&CK framework mapping. + +Your task is to extract Techniques, Tactics, and Procedures (TTPs) from documents. For each TTP found: +1. Identify the specific technique or method described +2. Categorize it (initial_access, execution, persistence, privilege_escalation, defense_evasion, credential_access, discovery, lateral_movement, collection, exfiltration, command_and_control, impact, reconnaissance, resource_development) +3. Note any attributed threat actor +4. Map to MITRE ATT&CK ID if possible +5. Include surrounding context for provenance + +Respond ONLY with a JSON object. No explanations outside JSON.""" + +TTP_EXTRACTION_PROMPT = """Extract all TTPs from this document. Return a JSON object with an "items" array. + +Each item must have: technique, category, description, actor (null if unknown), mitre_id (null if unknown), source_context (verbatim quote from document, max 200 chars). + + +{document} + + +Return JSON: +{{"items": [{{"technique": "...", "category": "...", "description": "...", "actor": null, "mitre_id": null, "source_context": "..."}}]}}""" + +# --- Tool Extraction --- + +TOOL_SYSTEM_PROMPT = """You are a cybersecurity researcher specializing in offensive tools, exploits, implants, and malware analysis. + +Your task is to extract named tools, exploits, implants, and malware from documents. For each tool found: +1. Identify the tool name and any aliases +2. Classify its capability (exploit, implant, framework, utility, spyware, rootkit, credential_theft, c2_framework, collection, obfuscation, proxy, trojan, weapon, platform) +3. Note target platforms and software +4. List any CVEs mentioned +5. Include surrounding context for provenance + +Respond ONLY with a JSON object. No explanations outside JSON.""" + +TOOL_EXTRACTION_PROMPT = """Extract all named tools, exploits, implants, and malware from this document. + +Known tools for reference (not exhaustive): {known_tools} + +Return a JSON object with an "items" array. Each item must have: name, aliases (list), description, capability, target_platforms (list), target_software (list), cves (list), mitre_techniques (list), source_context (verbatim quote, max 200 chars). + + +{document} + + +Return JSON: +{{"items": [{{"name": "...", "aliases": [], "description": "...", "capability": "...", "target_platforms": [], "target_software": [], "cves": [], "mitre_techniques": [], "source_context": "..."}}]}}""" + +# --- Tradecraft Extraction --- + +TRADECRAFT_SYSTEM_PROMPT = """You are an intelligence analyst specializing in tradecraft — the methods and techniques used in intelligence operations, espionage, and covert activities. + +Your task is to extract operational tradecraft from documents, including: +- HUMINT methods (recruitment, handling, dead drops, surveillance detection) +- SIGINT methods (interception, collection, processing) +- Cyber operations tradecraft (OPSEC, infrastructure setup, persistence techniques) +- Surveillance methods (physical and electronic) +- Communication security methods +- Counter-intelligence techniques + +For each method, note how it's actually used operationally and what countermeasures exist. + +Respond ONLY with a JSON object. No explanations outside JSON.""" + +TRADECRAFT_EXTRACTION_PROMPT = """Extract all tradecraft methods from this document. + +Return a JSON object with an "items" array. Each item must have: method, domain (humint/sigint/cyber/surveillance/commsec/counterintel), description, operational_notes (how it's actually used), countermeasures (how to detect/defend), source_context (verbatim quote, max 200 chars). + + +{document} + + +Return JSON: +{{"items": [{{"method": "...", "domain": "...", "description": "...", "operational_notes": "...", "countermeasures": "...", "source_context": "..."}}]}}""" + +# --- Surveillance Extraction --- + +SURVEILLANCE_SYSTEM_PROMPT = """You are an intelligence analyst specializing in surveillance technology and programs. + +Your task is to extract information about surveillance programs, intercept capabilities, collection methods, and vendor capabilities from documents. This includes: +- Government surveillance programs (PRISM, XKeyscore, etc.) +- Commercial surveillance tools (FinFisher, Pegasus, etc.) +- Interception capabilities and methods +- Data collection infrastructure +- Vendor/contractor capabilities + +Respond ONLY with a JSON object. No explanations outside JSON.""" + +SURVEILLANCE_EXTRACTION_PROMPT = """Extract all surveillance programs, capabilities, and methods from this document. + +Return a JSON object with an "items" array. Each item must have: program (or technique name), description, operator (government/vendor/unknown), capability_type (intercept/collection/analysis/targeting), mitre_id (null if unknown), source_context (verbatim quote, max 200 chars). + + +{document} + + +Return JSON: +{{"items": [{{"program": "...", "description": "...", "operator": null, "capability_type": "...", "mitre_id": null, "source_context": "..."}}]}}""" + +# --- Infrastructure Extraction --- + +INFRASTRUCTURE_SYSTEM_PROMPT = """You are a threat intelligence analyst specializing in identifying network infrastructure indicators. + +Your task is to extract infrastructure indicators from documents, including: +- IP addresses (IPv4 and IPv6) +- Domain names +- URLs +- File hashes (MD5, SHA1, SHA256) +- C2 server addresses +- Email addresses used in operations + +For each indicator, provide context about what it relates to and how it was used. +IMPORTANT: Only extract indicators that appear in the actual document text. Do not fabricate or guess indicators. + +Respond ONLY with a JSON object. No explanations outside JSON.""" + +INFRASTRUCTURE_EXTRACTION_PROMPT = """Extract all network infrastructure indicators from this document. + +Return a JSON object with an "items" array. Each item must have: indicator (the actual IP/domain/URL/hash), indicator_type (ipv4/ipv6/domain/url/hash/email), context (what this indicator relates to), source_context (verbatim surrounding text, max 200 chars). + + +{document} + + +Return JSON: +{{"items": [{{"indicator": "...", "indicator_type": "...", "context": "...", "source_context": "..."}}]}}""" + +# --- Entity Extraction --- + +ENTITY_SYSTEM_PROMPT = """You are an intelligence analyst specializing in entity identification and relationship mapping. + +Your task is to extract named entities from documents, including: +- People (officials, operators, researchers, targets) +- Organizations (agencies, companies, units, departments) +- Programs (codenames, project names, operation names) +- Codenames (tool names, operation names, program designators) +- Military/intelligence units + +For each entity, identify relationships to other entities mentioned in the same document. + +Respond ONLY with a JSON object. No explanations outside JSON.""" + +ENTITY_EXTRACTION_PROMPT = """Extract all named entities and their relationships from this document. + +Return a JSON object with an "items" array. Each item must have: name, type (person/org/program/codename/unit), description, relationships (list of {{"entity": "name", "relationship_type": "type"}}), source_context (verbatim quote, max 200 chars). + + +{document} + + +Return JSON: +{{"items": [{{"name": "...", "type": "...", "description": "...", "relationships": [{{"entity": "...", "relationship_type": "..."}}], "source_context": "..."}}]}}""" + +# --- Interactive Query Prompts --- + +QUERY_SYSTEM_PROMPT = """You are an intelligence analyst with access to a corpus of documents. Answer questions based on the provided document excerpts. Always cite your sources by referencing the document title or ID. + +If the provided context doesn't contain enough information to answer confidently, say so explicitly rather than speculating.""" + +QUERY_PROMPT = """Based on the following document excerpts, answer the user's question. + +{context} + +Question: {question} + +Provide a detailed answer with citations to specific documents.""" + +EMULATE_PROMPT = """Based on the extracted intelligence data below, create a detailed TTP playbook for emulating the threat actor: {actor} + +Include: +1. Initial access methods they've used +2. Persistence mechanisms +3. Lateral movement techniques +4. Data collection methods +5. Exfiltration techniques +6. C2 infrastructure patterns +7. OPSEC practices +8. Recommended tools for emulation + +{context} + +Generate the playbook:""" + +COMPARE_PROMPT = """Compare these two tools/capabilities based on the extracted intelligence: + +Tool A: {tool_a} +Tool B: {tool_b} + +{context} + +Compare across: capability, target platforms, sophistication, attribution, detection methods.""" + +BRIEF_PROMPT = """Generate an intelligence briefing on the topic: {topic} + +Based on the following extracted intelligence: +{context} + +Structure the briefing with: +1. Executive Summary +2. Key Findings +3. Technical Details +4. Threat Assessment +5. Recommended Actions""" + +SCENARIO_PROMPT = """Generate an adversary simulation scenario targeting: {target} + +Based on the following threat intelligence: +{context} + +Include: +1. Threat actor profile (based on extracted TTPs) +2. Attack chain (step by step) +3. Tools and infrastructure needed +4. Detection opportunities for blue team +5. Recommended mitigations""" diff --git a/extractors/__init__.py b/extractors/__init__.py index e69de29..9824503 100644 --- a/extractors/__init__.py +++ b/extractors/__init__.py @@ -0,0 +1,8 @@ +from .base import BaseExtractor +from .ttp_extractor import TTPExtractor +from .tool_extractor import ToolExtractor +from .tradecraft import TradecraftExtractor +from .surveillance import SurveillanceExtractor +from .infrastructure import InfrastructureExtractor +from .entity_extractor import EntityExtractor +from .mitre_mapper import MITREMapper diff --git a/extractors/base.py b/extractors/base.py new file mode 100644 index 0000000..9c676cc --- /dev/null +++ b/extractors/base.py @@ -0,0 +1,236 @@ +"""Base extractor for Claude-powered intelligence extraction.""" +import json +import logging +from abc import ABC, abstractmethod +from datetime import datetime +from typing import Optional + +import anthropic + +from storage.db import MosaicDB +from storage.models import TokenUsage +from utils.api_limiter import APIRateLimiter, CircuitBreakerOpen +from utils.text_utils import chunk_text, estimate_tokens + +logger = logging.getLogger(__name__) + + +class BaseExtractor(ABC): + """Abstract base for all Claude-powered extractors.""" + + EXTRACTOR_NAME = "base" # Override in subclasses + + def __init__(self, db: MosaicDB, api_key: str, + model: str = "claude-sonnet-4-20250514", + max_tokens: int = 4096, + rate_limiter: Optional[APIRateLimiter] = None): + self.db = db + self.client = anthropic.Anthropic(api_key=api_key) + self.model = model + self.max_tokens = max_tokens + self.rate_limiter = rate_limiter or APIRateLimiter() + + @abstractmethod + def get_system_prompt(self) -> str: + """Return the system prompt for this extractor.""" + pass + + @abstractmethod + def get_extraction_prompt(self, text: str) -> str: + """Return the user prompt with document content wrapped in XML tags.""" + pass + + @abstractmethod + def get_json_schema(self) -> dict: + """Return the expected JSON schema for validation.""" + pass + + @abstractmethod + def process_response(self, data: dict, doc_id: str) -> int: + """Process validated extraction response and store results. + Returns count of items extracted.""" + pass + + def extract_from_document(self, doc_id: str, text: str, doc_type: str = "text", + source: str = "", metadata: Optional[dict] = None) -> int: + """Extract intelligence from a single document. + + Handles chunking, API calls, validation, and storage. + Returns total count of items extracted. + """ + # Set status to IN_PROGRESS + self.db.set_extraction_status(doc_id, self.EXTRACTOR_NAME, "IN_PROGRESS") + + try: + # Chunk the document + chunks = chunk_text(text, doc_type=doc_type, metadata=metadata) + total_extracted = 0 + total_input_tokens = 0 + total_output_tokens = 0 + + for chunk in chunks: + try: + # Rate limit + self.rate_limiter.wait_if_needed() + + # Build prompt with XML-wrapped content + user_prompt = self.get_extraction_prompt(chunk['text']) + + # Call Claude + response = self.client.messages.create( + model=self.model, + max_tokens=self.max_tokens, + system=self.get_system_prompt(), + messages=[{"role": "user", "content": user_prompt}] + ) + + # Track tokens + input_tokens = response.usage.input_tokens + output_tokens = response.usage.output_tokens + total_input_tokens += input_tokens + total_output_tokens += output_tokens + + self.rate_limiter.record_success() + + # Extract text from response + response_text = response.content[0].text.strip() + + # Parse JSON from response + data = self._parse_json_response(response_text) + if data is None: + logger.warning("Failed to parse JSON from response for doc %s chunk %d", + doc_id, chunk.get('chunk_index', 0)) + continue + + # Validate against schema + if not self._validate_response(data): + logger.warning("Schema validation failed for doc %s chunk %d", + doc_id, chunk.get('chunk_index', 0)) + continue + + # Process and store + count = self.process_response(data, doc_id) + total_extracted += count + + except CircuitBreakerOpen: + logger.error("Circuit breaker open — halting extraction") + self.db.set_extraction_status( + doc_id, self.EXTRACTOR_NAME, "FAILED", + error="Circuit breaker open" + ) + return total_extracted + except anthropic.APIError as e: + self.rate_limiter.record_failure(str(e)) + logger.error("API error on doc %s: %s", doc_id, e) + continue + except Exception as e: + logger.error("Extraction error on doc %s chunk %d: %s", + doc_id, chunk.get('chunk_index', 0), e) + continue + + # Log token usage + self.db.log_token_usage(TokenUsage( + timestamp=datetime.utcnow().isoformat(), + source=source, + doc_id=doc_id, + extractor=self.EXTRACTOR_NAME, + input_tokens=total_input_tokens, + output_tokens=total_output_tokens, + model=self.model, + )) + + # Update status + self.db.set_extraction_status( + doc_id, self.EXTRACTOR_NAME, "DONE", + token_count=total_input_tokens + total_output_tokens + ) + self.db.update_document_extraction_time(doc_id) + + return total_extracted + + except Exception as e: + logger.error("Extraction failed for doc %s: %s", doc_id, e) + self.db.set_extraction_status( + doc_id, self.EXTRACTOR_NAME, "FAILED", error=str(e) + ) + return 0 + + def estimate_tokens(self, text: str, doc_type: str = "text", + metadata: Optional[dict] = None) -> dict: + """Estimate token cost without calling the API (dry-run).""" + chunks = chunk_text(text, doc_type=doc_type, metadata=metadata) + total_input = 0 + system_tokens = estimate_tokens(self.get_system_prompt()) + + for chunk in chunks: + prompt = self.get_extraction_prompt(chunk['text']) + total_input += estimate_tokens(prompt) + system_tokens + + return { + 'chunks': len(chunks), + 'estimated_input_tokens': total_input, + 'estimated_output_tokens': len(chunks) * 2000, # rough estimate + } + + def _parse_json_response(self, text: str) -> Optional[dict]: + """Parse JSON from Claude response, handling markdown code blocks.""" + # Try direct parse + try: + return json.loads(text) + except json.JSONDecodeError: + pass + + # Try extracting from markdown code block + import re + match = re.search(r'```(?:json)?\s*\n(.*?)\n```', text, re.DOTALL) + if match: + try: + return json.loads(match.group(1)) + except json.JSONDecodeError: + pass + + # Try finding JSON object/array in text + for start_char, end_char in [('{', '}'), ('[', ']')]: + start = text.find(start_char) + if start >= 0: + # Find matching end + depth = 0 + for i in range(start, len(text)): + if text[i] == start_char: + depth += 1 + elif text[i] == end_char: + depth -= 1 + if depth == 0: + try: + return json.loads(text[start:i+1]) + except json.JSONDecodeError: + break + + logger.warning("Could not parse JSON from response: %s...", text[:200]) + return None + + def _validate_response(self, data: dict) -> bool: + """Validate response against expected schema. Reject unexpected fields.""" + schema = self.get_json_schema() + if not schema: + return True + + required_keys = schema.get('required', []) + allowed_keys = set(schema.get('properties', {}).keys()) + + if isinstance(data, dict): + # Check for unexpected top-level keys + if 'items' in schema.get('properties', {}): + # Wrapper object with 'items' array + if 'items' not in data: + logger.warning("Response missing 'items' key") + return False + unexpected = set(data.keys()) - allowed_keys + if unexpected: + logger.warning("Unexpected fields in response: %s", unexpected) + return False + return True + elif isinstance(data, list): + return True + + return False diff --git a/extractors/entity_extractor.py b/extractors/entity_extractor.py new file mode 100644 index 0000000..d0ceb84 --- /dev/null +++ b/extractors/entity_extractor.py @@ -0,0 +1,42 @@ +"""Entity extractor (people, orgs, programs, codenames, units).""" +import logging + +from .base import BaseExtractor +from storage.models import Entity +from analysis.prompts import ENTITY_SYSTEM_PROMPT, ENTITY_EXTRACTION_PROMPT + +logger = logging.getLogger(__name__) + + +class EntityExtractor(BaseExtractor): + EXTRACTOR_NAME = "entity" + + def get_system_prompt(self) -> str: + return ENTITY_SYSTEM_PROMPT + + def get_extraction_prompt(self, text: str) -> str: + return ENTITY_EXTRACTION_PROMPT.format(document=text) + + def get_json_schema(self) -> dict: + return {"properties": {"items": {"type": "array"}}, "required": ["items"]} + + def process_response(self, data: dict, doc_id: str) -> int: + items = data.get('items', []) + if isinstance(data, list): + items = data + count = 0 + for item in items: + try: + entity = Entity( + name=item.get('name', ''), + type=item.get('type', ''), + description=item.get('description', ''), + relationships=item.get('relationships', []), + source_doc_ids=[doc_id], + source_context=item.get('source_context', '')[:500], + ) + self.db.insert_entity(entity) + count += 1 + except Exception as e: + logger.warning("Failed to process entity item: %s", e) + return count diff --git a/extractors/infrastructure.py b/extractors/infrastructure.py new file mode 100644 index 0000000..665bd1f --- /dev/null +++ b/extractors/infrastructure.py @@ -0,0 +1,51 @@ +"""Infrastructure indicator extractor for Mosaic (IPs, domains, C2, hashes).""" +import logging + +from .base import BaseExtractor +from storage.models import Infrastructure +from analysis.prompts import INFRASTRUCTURE_SYSTEM_PROMPT, INFRASTRUCTURE_EXTRACTION_PROMPT + +logger = logging.getLogger(__name__) + + +class InfrastructureExtractor(BaseExtractor): + EXTRACTOR_NAME = "infrastructure" + + def __init__(self, *args, collection: str = "", doc_title: str = "", + doc_date: str = None, **kwargs): + super().__init__(*args, **kwargs) + self.collection = collection + self.doc_title = doc_title + self.doc_date = doc_date + + def get_system_prompt(self) -> str: + return INFRASTRUCTURE_SYSTEM_PROMPT + + def get_extraction_prompt(self, text: str) -> str: + return INFRASTRUCTURE_EXTRACTION_PROMPT.format(document=text) + + def get_json_schema(self) -> dict: + return {"properties": {"items": {"type": "array"}}, "required": ["items"]} + + def process_response(self, data: dict, doc_id: str) -> int: + items = data.get('items', []) + if isinstance(data, list): + items = data + count = 0 + for item in items: + try: + infra = Infrastructure( + indicator=item.get('indicator', ''), + indicator_type=item.get('indicator_type', ''), + context=item.get('context', ''), + source_doc_ids=[doc_id], + source_context=item.get('source_context', '')[:500], + collection=self.collection, + doc_title=self.doc_title, + doc_date=self.doc_date, + ) + self.db.insert_infrastructure(infra) + count += 1 + except Exception as e: + logger.warning("Failed to process infrastructure item: %s", e) + return count diff --git a/extractors/mitre_mapper.py b/extractors/mitre_mapper.py new file mode 100644 index 0000000..4811983 --- /dev/null +++ b/extractors/mitre_mapper.py @@ -0,0 +1,131 @@ +"""MITRE ATT&CK mapper — maps extracted TTPs to ATT&CK technique IDs.""" +import json +import hashlib +import logging +from pathlib import Path +from typing import Optional + +logger = logging.getLogger(__name__) + +ATTACK_DATA_PATH = Path(__file__).parent.parent / 'data' / 'mitre_attack.json' + + +class MITREMapper: + """Map extracted TTPs to MITRE ATT&CK technique IDs.""" + + def __init__(self): + self.techniques = {} + self.tactics = {} + self._loaded = False + self._load_attack_data() + + def _load_attack_data(self): + """Load ATT&CK data and verify integrity.""" + if not ATTACK_DATA_PATH.exists(): + logger.warning("MITRE ATT&CK data not found at %s", ATTACK_DATA_PATH) + return + + try: + with open(ATTACK_DATA_PATH) as f: + data = json.load(f) + + # Load tactics + for tactic in data.get('tactics', []): + tid = tactic['id'] + self.tactics[tid] = tactic['name'] + self.tactics[tactic['name'].lower()] = tid + + # Load techniques (if populated) + for technique in data.get('techniques', []): + tid = technique.get('id', '') + name = technique.get('name', '') + if tid and name: + self.techniques[tid] = { + 'name': name, + 'tactic': technique.get('tactic', ''), + 'description': technique.get('description', ''), + } + self.techniques[name.lower()] = tid + + self._loaded = True + logger.debug("Loaded %d tactics, %d techniques", + len(data.get('tactics', [])), + len(data.get('techniques', []))) + + except Exception as e: + logger.error("Failed to load ATT&CK data: %s", e) + + def map_technique(self, technique_name: str, category: str = "") -> Optional[str]: + """Map a technique name or description to a MITRE ATT&CK ID. + + Returns ATT&CK ID (e.g., 'T1566') or None if no match. + """ + if not self._loaded: + return None + + # Direct ID match + if technique_name.upper().startswith('T') and technique_name[1:].replace('.', '').isdigit(): + if technique_name.upper() in self.techniques: + return technique_name.upper() + + # Name match + lower_name = technique_name.lower().strip() + if lower_name in self.techniques: + result = self.techniques[lower_name] + if isinstance(result, str): + return result + + # Category to tactic mapping + category_tactic_map = { + 'initial_access': 'TA0001', + 'execution': 'TA0002', + 'persistence': 'TA0003', + 'privilege_escalation': 'TA0004', + 'defense_evasion': 'TA0005', + 'credential_access': 'TA0006', + 'discovery': 'TA0007', + 'lateral_movement': 'TA0008', + 'collection': 'TA0009', + 'exfiltration': 'TA0010', + 'command_and_control': 'TA0011', + 'impact': 'TA0040', + 'resource_development': 'TA0042', + 'reconnaissance': 'TA0043', + # Aliases + 'c2': 'TA0011', + 'exfil': 'TA0010', + 'privesc': 'TA0004', + 'recon': 'TA0043', + 'surveillance': 'TA0009', + } + + return category_tactic_map.get(category.lower().strip()) + + def get_tactic_name(self, tactic_id: str) -> Optional[str]: + """Get tactic name from ID.""" + return self.tactics.get(tactic_id) + + def get_technique_info(self, technique_id: str) -> Optional[dict]: + """Get technique details from ID.""" + return self.techniques.get(technique_id) if isinstance( + self.techniques.get(technique_id), dict) else None + + def enrich_ttp(self, technique: str, category: str) -> dict: + """Enrich a TTP with MITRE mapping. + + Returns dict with mitre_id and tactic if found. + """ + result = {'mitre_id': None, 'tactic': None} + + mitre_id = self.map_technique(technique, category) + if mitre_id: + if mitre_id.startswith('TA'): + result['tactic'] = mitre_id + result['tactic_name'] = self.get_tactic_name(mitre_id) + else: + result['mitre_id'] = mitre_id + info = self.get_technique_info(mitre_id) + if info: + result['tactic'] = info.get('tactic') + + return result diff --git a/extractors/surveillance.py b/extractors/surveillance.py new file mode 100644 index 0000000..305a191 --- /dev/null +++ b/extractors/surveillance.py @@ -0,0 +1,45 @@ +"""Surveillance program/capability extractor for Mosaic.""" +import logging + +from .base import BaseExtractor +from storage.models import TTP +from analysis.prompts import SURVEILLANCE_SYSTEM_PROMPT, SURVEILLANCE_EXTRACTION_PROMPT + +logger = logging.getLogger(__name__) + + +class SurveillanceExtractor(BaseExtractor): + EXTRACTOR_NAME = "surveillance" + + def get_system_prompt(self) -> str: + return SURVEILLANCE_SYSTEM_PROMPT + + def get_extraction_prompt(self, text: str) -> str: + return SURVEILLANCE_EXTRACTION_PROMPT.format(document=text) + + def get_json_schema(self) -> dict: + return {"properties": {"items": {"type": "array"}}, "required": ["items"]} + + def process_response(self, data: dict, doc_id: str) -> int: + items = data.get('items', []) + if isinstance(data, list): + items = data + count = 0 + for item in items: + try: + # Store surveillance as TTPs with surveillance category + ttp = TTP( + technique=item.get('program', item.get('technique', '')), + category='surveillance', + description=item.get('description', ''), + actor=item.get('operator', item.get('actor', None)), + mitre_id=item.get('mitre_id'), + source_doc_ids=[doc_id], + source_context=item.get('source_context', '')[:500], + ) + ttp.compute_dedup_key() + if self.db.insert_ttp(ttp): + count += 1 + except Exception as e: + logger.warning("Failed to process surveillance item: %s", e) + return count diff --git a/extractors/tool_extractor.py b/extractors/tool_extractor.py new file mode 100644 index 0000000..53df4ae --- /dev/null +++ b/extractors/tool_extractor.py @@ -0,0 +1,78 @@ +"""Named tool/exploit/implant extractor for Mosaic.""" +import json +import logging +from pathlib import Path +from typing import Optional + +from .base import BaseExtractor +from storage.db import MosaicDB +from storage.models import ExtractedTool +from analysis.prompts import TOOL_SYSTEM_PROMPT, TOOL_EXTRACTION_PROMPT + +logger = logging.getLogger(__name__) + +KNOWN_TOOLS_PATH = Path(__file__).parent.parent / 'data' / 'known_tools.json' + + +class ToolExtractor(BaseExtractor): + EXTRACTOR_NAME = "tool" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.known_tools = self._load_known_tools() + + def _load_known_tools(self) -> list[str]: + try: + with open(KNOWN_TOOLS_PATH) as f: + data = json.load(f) + names = [] + for tool in data.get('tools', []): + names.append(tool['name'].lower()) + for alias in tool.get('aliases', []): + names.append(alias.lower()) + return names + except Exception as e: + logger.warning("Could not load known tools: %s", e) + return [] + + def get_system_prompt(self) -> str: + return TOOL_SYSTEM_PROMPT + + def get_extraction_prompt(self, text: str) -> str: + # Include known tools as context + known_hint = ', '.join(self.known_tools[:50]) if self.known_tools else 'none loaded' + return TOOL_EXTRACTION_PROMPT.format(document=text, known_tools=known_hint) + + def get_json_schema(self) -> dict: + return { + "properties": { + "items": {"type": "array"}, + }, + "required": ["items"] + } + + def process_response(self, data: dict, doc_id: str) -> int: + items = data.get('items', []) + if isinstance(data, list): + items = data + count = 0 + for item in items: + try: + tool = ExtractedTool( + name=item.get('name', ''), + aliases=item.get('aliases', []), + description=item.get('description', ''), + capability=item.get('capability', ''), + target_platforms=item.get('target_platforms', []), + target_software=item.get('target_software', []), + cves=item.get('cves', []), + source_doc_ids=[doc_id], + source_context=item.get('source_context', '')[:500], + mitre_techniques=item.get('mitre_techniques', []), + ) + tool.compute_dedup_key() + if self.db.insert_extracted_tool(tool): + count += 1 + except Exception as e: + logger.warning("Failed to process tool item: %s", e) + return count diff --git a/extractors/tradecraft.py b/extractors/tradecraft.py new file mode 100644 index 0000000..65fae5b --- /dev/null +++ b/extractors/tradecraft.py @@ -0,0 +1,43 @@ +"""Tradecraft extractor (HUMINT/SIGINT methods, OPSEC procedures).""" +import logging + +from .base import BaseExtractor +from storage.models import Tradecraft +from analysis.prompts import TRADECRAFT_SYSTEM_PROMPT, TRADECRAFT_EXTRACTION_PROMPT + +logger = logging.getLogger(__name__) + + +class TradecraftExtractor(BaseExtractor): + EXTRACTOR_NAME = "tradecraft" + + def get_system_prompt(self) -> str: + return TRADECRAFT_SYSTEM_PROMPT + + def get_extraction_prompt(self, text: str) -> str: + return TRADECRAFT_EXTRACTION_PROMPT.format(document=text) + + def get_json_schema(self) -> dict: + return {"properties": {"items": {"type": "array"}}, "required": ["items"]} + + def process_response(self, data: dict, doc_id: str) -> int: + items = data.get('items', []) + if isinstance(data, list): + items = data + count = 0 + for item in items: + try: + tc = Tradecraft( + method=item.get('method', ''), + domain=item.get('domain', ''), + description=item.get('description', ''), + operational_notes=item.get('operational_notes', ''), + countermeasures=item.get('countermeasures', ''), + source_doc_ids=[doc_id], + source_context=item.get('source_context', '')[:500], + ) + self.db.insert_tradecraft(tc) + count += 1 + except Exception as e: + logger.warning("Failed to process tradecraft item: %s", e) + return count diff --git a/extractors/ttp_extractor.py b/extractors/ttp_extractor.py new file mode 100644 index 0000000..5f7d2be --- /dev/null +++ b/extractors/ttp_extractor.py @@ -0,0 +1,51 @@ +"""TTP (Techniques, Tactics, Procedures) extractor for Mosaic.""" +import logging +from typing import Optional + +from .base import BaseExtractor +from storage.db import MosaicDB +from storage.models import TTP +from analysis.prompts import TTP_SYSTEM_PROMPT, TTP_EXTRACTION_PROMPT + +logger = logging.getLogger(__name__) + + +class TTPExtractor(BaseExtractor): + EXTRACTOR_NAME = "ttp" + + def get_system_prompt(self) -> str: + return TTP_SYSTEM_PROMPT + + def get_extraction_prompt(self, text: str) -> str: + return TTP_EXTRACTION_PROMPT.format(document=text) + + def get_json_schema(self) -> dict: + return { + "properties": { + "items": {"type": "array"}, + }, + "required": ["items"] + } + + def process_response(self, data: dict, doc_id: str) -> int: + items = data.get('items', []) + if isinstance(data, list): + items = data + count = 0 + for item in items: + try: + ttp = TTP( + technique=item.get('technique', ''), + category=item.get('category', ''), + description=item.get('description', ''), + actor=item.get('actor'), + mitre_id=item.get('mitre_id'), + source_doc_ids=[doc_id], + source_context=item.get('source_context', '')[:500], + ) + ttp.compute_dedup_key() + if self.db.insert_ttp(ttp): + count += 1 + except Exception as e: + logger.warning("Failed to process TTP item: %s", e) + return count