commit aa0f35803f0ad77ce54c1c008d1b8fd8be10e3ef Author: n0mad1k Date: Thu Mar 19 08:34:32 2026 -0400 Initial scaffold: .gitignore, config, requirements, data seeds, directory structure diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..fc494ee --- /dev/null +++ b/.gitignore @@ -0,0 +1,9 @@ +config/mosaic.conf +output/ +*.db +*.db-wal +*.db-shm +.env +__pycache__/ +.venv/ +*.pyc diff --git a/analysis/__init__.py b/analysis/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/collectors/__init__.py b/collectors/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/collectors/profiles/cablegate.yaml b/collectors/profiles/cablegate.yaml new file mode 100644 index 0000000..cc19050 --- /dev/null +++ b/collectors/profiles/cablegate.yaml @@ -0,0 +1,20 @@ +name: cablegate +description: "US Diplomatic Cables (Cablegate)" +base_url: "https://wikileaks.org/plusd/" +domain_allowlist: + - wikileaks.org +crawl_rules: + max_depth: 2 + max_documents: 300000 + follow_links: true + url_patterns: + - "/plusd/cables/*" + exclude_patterns: + - "*.css" + - "*.js" + - "*.ico" +content_types: + - text/html +expected_min_count: 250000 +rate_limit: 0.5 +mirrors: [] diff --git a/collectors/profiles/emails.yaml b/collectors/profiles/emails.yaml new file mode 100644 index 0000000..64aa058 --- /dev/null +++ b/collectors/profiles/emails.yaml @@ -0,0 +1,21 @@ +name: emails +description: "Email Dumps (DNC, Podesta, etc.)" +base_url: "https://wikileaks.org/dnc-emails/" +domain_allowlist: + - wikileaks.org +crawl_rules: + max_depth: 2 + max_documents: 50000 + follow_links: true + url_patterns: + - "/dnc-emails/emailid/*" + - "/podesta-emails/emailid/*" + exclude_patterns: + - "*.css" + - "*.js" + - "*.ico" +content_types: + - text/html +expected_min_count: 20000 +rate_limit: 0.5 +mirrors: [] diff --git a/collectors/profiles/example.yaml b/collectors/profiles/example.yaml new file mode 100644 index 0000000..3fbf3ea --- /dev/null +++ b/collectors/profiles/example.yaml @@ -0,0 +1,27 @@ +# Example source profile — copy and customize for your own sources +# Place in collectors/profiles/ or pass with --profile flag +name: my_source +description: "Description of the document source" +base_url: "https://example.com/documents/" +domain_allowlist: + - example.com + - www.example.com +crawl_rules: + max_depth: 2 + max_documents: 500 + follow_links: true + url_patterns: + - "/documents/*" + - "*.pdf" + exclude_patterns: + - "*.css" + - "*.js" + - "*.ico" + - "*.png" + - "*.jpg" +content_types: + - text/html + - application/pdf +expected_min_count: 100 +rate_limit: 1.0 +mirrors: [] diff --git a/collectors/profiles/spyfiles.yaml b/collectors/profiles/spyfiles.yaml new file mode 100644 index 0000000..0b2a0bb --- /dev/null +++ b/collectors/profiles/spyfiles.yaml @@ -0,0 +1,23 @@ +name: spyfiles +description: "Surveillance Industry Documents (Spy Files)" +base_url: "https://wikileaks.org/spyfiles/" +domain_allowlist: + - wikileaks.org + - file.wikileaks.org +crawl_rules: + max_depth: 3 + max_documents: 2000 + follow_links: true + url_patterns: + - "/spyfiles/document/*" + - "/spyfiles/files/*" + exclude_patterns: + - "*.css" + - "*.js" + - "*.ico" +content_types: + - text/html + - application/pdf +expected_min_count: 200 +rate_limit: 1.0 +mirrors: [] diff --git a/collectors/profiles/vault7.yaml b/collectors/profiles/vault7.yaml new file mode 100644 index 0000000..1f0ac52 --- /dev/null +++ b/collectors/profiles/vault7.yaml @@ -0,0 +1,27 @@ +name: vault7 +description: "CIA Hacking Tools (Vault 7)" +base_url: "https://wikileaks.org/vault7/" +domain_allowlist: + - wikileaks.org + - file.wikileaks.org +crawl_rules: + max_depth: 3 + max_documents: 5000 + follow_links: true + url_patterns: + - "/vault7/#*" + - "/vault7/document/*" + exclude_patterns: + - "*.css" + - "*.js" + - "*.ico" + - "*.png" + - "*.jpg" + - "*.gif" +content_types: + - text/html + - application/pdf +expected_min_count: 800 +rate_limit: 1.0 +mirrors: + - "https://file.wikileaks.org/vault7/" diff --git a/collectors/profiles/vault8.yaml b/collectors/profiles/vault8.yaml new file mode 100644 index 0000000..4a190fd --- /dev/null +++ b/collectors/profiles/vault8.yaml @@ -0,0 +1,27 @@ +name: vault8 +description: "CIA Source Code (Vault 8)" +base_url: "https://wikileaks.org/vault8/" +domain_allowlist: + - wikileaks.org + - file.wikileaks.org +crawl_rules: + max_depth: 3 + max_documents: 1000 + follow_links: true + url_patterns: + - "/vault8/#*" + - "/vault8/document/*" + exclude_patterns: + - "*.css" + - "*.js" + - "*.ico" + - "*.png" + - "*.jpg" +content_types: + - text/html + - application/pdf + - text/plain +expected_min_count: 100 +rate_limit: 1.0 +mirrors: + - "https://file.wikileaks.org/vault8/" diff --git a/data/known_tools.json b/data/known_tools.json new file mode 100644 index 0000000..6f86776 --- /dev/null +++ b/data/known_tools.json @@ -0,0 +1,44 @@ +{ + "version": "1.0", + "tools": [ + {"name": "HIVE", "aliases": ["Hive"], "category": "implant"}, + {"name": "Grasshopper", "aliases": [], "category": "framework"}, + {"name": "Marble", "aliases": ["Marble Framework"], "category": "obfuscation"}, + {"name": "Weeping Angel", "aliases": ["WA"], "category": "implant"}, + {"name": "Dark Matter", "aliases": [], "category": "implant"}, + {"name": "Athena", "aliases": [], "category": "implant"}, + {"name": "AfterMidnight", "aliases": ["After Midnight"], "category": "implant"}, + {"name": "Brutal Kangaroo", "aliases": ["BK"], "category": "framework"}, + {"name": "OutlawCountry", "aliases": ["Outlaw Country"], "category": "implant"}, + {"name": "BothanSpy", "aliases": ["Bothan Spy"], "category": "credential_theft"}, + {"name": "Gyrfalcon", "aliases": [], "category": "credential_theft"}, + {"name": "HighRise", "aliases": ["High Rise"], "category": "proxy"}, + {"name": "Achilles", "aliases": [], "category": "trojan"}, + {"name": "Aeris", "aliases": [], "category": "implant"}, + {"name": "SeaPea", "aliases": ["Sea Pea"], "category": "rootkit"}, + {"name": "CouchPotato", "aliases": ["Couch Potato"], "category": "collection"}, + {"name": "Dumbo", "aliases": [], "category": "evasion"}, + {"name": "Pandemic", "aliases": [], "category": "implant"}, + {"name": "Cherry Blossom", "aliases": ["CherryBlossom", "CB"], "category": "router_exploit"}, + {"name": "McNugget", "aliases": [], "category": "implant"}, + {"name": "FinFisher", "aliases": ["FinSpy"], "category": "spyware"}, + {"name": "Pegasus", "aliases": [], "category": "spyware"}, + {"name": "Predator", "aliases": [], "category": "spyware"}, + {"name": "WARRIOR PRIDE", "aliases": [], "category": "mobile_implant"}, + {"name": "PRISM", "aliases": [], "category": "collection_program"}, + {"name": "XKeyscore", "aliases": ["XKEYSCORE", "XKS"], "category": "sigint"}, + {"name": "TURBINE", "aliases": [], "category": "implant_management"}, + {"name": "FOXACID", "aliases": ["FoxAcid"], "category": "exploit_server"}, + {"name": "QUANTUM", "aliases": ["QuantumInsert"], "category": "injection"}, + {"name": "Regin", "aliases": [], "category": "platform"}, + {"name": "Flame", "aliases": ["Flamer", "sKyWIper"], "category": "platform"}, + {"name": "Duqu", "aliases": ["Duqu 2.0"], "category": "platform"}, + {"name": "Stuxnet", "aliases": [], "category": "weapon"}, + {"name": "EternalBlue", "aliases": ["MS17-010"], "category": "exploit"}, + {"name": "DoublePulsar", "aliases": [], "category": "implant"}, + {"name": "EternalRomance", "aliases": [], "category": "exploit"}, + {"name": "Cobalt Strike", "aliases": ["CS", "CobaltStrike"], "category": "c2_framework"}, + {"name": "Metasploit", "aliases": ["MSF"], "category": "framework"}, + {"name": "Mimikatz", "aliases": [], "category": "credential_theft"} + ] +} diff --git a/data/mitre_attack.json b/data/mitre_attack.json new file mode 100644 index 0000000..2e287fa --- /dev/null +++ b/data/mitre_attack.json @@ -0,0 +1,22 @@ +{ + "version": "15.1", + "sha256": "placeholder_will_be_updated_on_first_download", + "source": "https://raw.githubusercontent.com/mitre/cti/master/enterprise-attack/enterprise-attack.json", + "tactics": [ + {"id": "TA0001", "name": "Initial Access"}, + {"id": "TA0002", "name": "Execution"}, + {"id": "TA0003", "name": "Persistence"}, + {"id": "TA0004", "name": "Privilege Escalation"}, + {"id": "TA0005", "name": "Defense Evasion"}, + {"id": "TA0006", "name": "Credential Access"}, + {"id": "TA0007", "name": "Discovery"}, + {"id": "TA0008", "name": "Lateral Movement"}, + {"id": "TA0009", "name": "Collection"}, + {"id": "TA0010", "name": "Exfiltration"}, + {"id": "TA0011", "name": "Command and Control"}, + {"id": "TA0040", "name": "Impact"}, + {"id": "TA0042", "name": "Resource Development"}, + {"id": "TA0043", "name": "Reconnaissance"} + ], + "techniques": [] +} diff --git a/extractors/__init__.py b/extractors/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/parsers/__init__.py b/parsers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..bff3564 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,8 @@ +click>=8.1.7,<9 +rich>=13.7.0,<14 +anthropic>=0.40.0,<1 +httpx>=0.27.0,<1 +beautifulsoup4>=4.12.3,<5 +PyMuPDF>=1.24.0,<2 +python-magic>=0.4.27,<1 +pyyaml>=6.0.1,<7 diff --git a/storage/__init__.py b/storage/__init__.py new file mode 100644 index 0000000..73c160b --- /dev/null +++ b/storage/__init__.py @@ -0,0 +1,4 @@ +from .db import MosaicDB +from .models import Document, ExtractedTool, TTP, Tradecraft, Entity, Infrastructure +from .migrations import MigrationRunner +from .jsonl_store import JSONLStore diff --git a/storage/db.py b/storage/db.py new file mode 100644 index 0000000..d23d67a --- /dev/null +++ b/storage/db.py @@ -0,0 +1,447 @@ +"""SQLite database interface for Mosaic with FTS5 and WAL mode.""" +import sqlite3 +import logging +import re +from pathlib import Path +from typing import Optional +from datetime import datetime + +from .models import ( + Document, ExtractedTool, TTP, Tradecraft, Entity, + Infrastructure, ExtractionStatus, TokenUsage +) +from .migrations import MigrationRunner + +logger = logging.getLogger(__name__) + +# FTS5 special characters that need escaping +FTS5_SPECIAL = re.compile(r'["\*\(\)\-\:]') +MAX_FTS_QUERY_LENGTH = 500 + + +class MosaicDB: + """SQLite database with FTS5 for Mosaic intelligence storage.""" + + def __init__(self, db_path: str = "mosaic.db"): + self.db_path = db_path + self.conn = self._connect() + self._setup() + + def _connect(self) -> sqlite3.Connection: + conn = sqlite3.connect(self.db_path) + conn.row_factory = sqlite3.Row + conn.execute("PRAGMA journal_mode=WAL") + conn.execute("PRAGMA foreign_keys=ON") + conn.execute("PRAGMA busy_timeout=5000") + return conn + + def _setup(self): + runner = MigrationRunner(self.conn) + runner.run() + # Crash recovery: reset any IN_PROGRESS states + self._reset_in_progress() + + def _reset_in_progress(self): + """Reset stale IN_PROGRESS/DOWNLOADING states on startup (crash recovery).""" + self.conn.execute( + "UPDATE documents SET status = 'PENDING' WHERE status = 'DOWNLOADING'" + ) + self.conn.execute( + "UPDATE documents SET parse_status = 'PENDING' WHERE parse_status = 'IN_PROGRESS'" + ) + self.conn.execute( + "UPDATE extraction_status SET status = 'PENDING' WHERE status = 'IN_PROGRESS'" + ) + self.conn.commit() + logger.debug("Reset stale in-progress states (crash recovery)") + + def close(self): + if self.conn: + self.conn.close() + + # --- Document operations --- + + def insert_document(self, doc: Document) -> bool: + """Insert or update a document. Returns True if new, False if updated.""" + existing = self.get_document(doc.id) + d = doc.to_dict() + if existing: + cols = ', '.join(f"{k} = ?" for k in d if k != 'id') + vals = [v for k, v in d.items() if k != 'id'] + [doc.id] + self.conn.execute(f"UPDATE documents SET {cols} WHERE id = ?", vals) + self.conn.commit() + return False + else: + cols = ', '.join(d.keys()) + placeholders = ', '.join('?' * len(d)) + self.conn.execute( + f"INSERT INTO documents ({cols}) VALUES ({placeholders})", + list(d.values()) + ) + self.conn.commit() + return True + + def get_document(self, doc_id: str) -> Optional[Document]: + cursor = self.conn.execute( + "SELECT * FROM documents WHERE id = ?", (doc_id,) + ) + row = cursor.fetchone() + return Document.from_row(row) if row else None + + def get_document_by_url(self, url: str) -> Optional[Document]: + cursor = self.conn.execute( + "SELECT * FROM documents WHERE source_url = ?", (url,) + ) + row = cursor.fetchone() + return Document.from_row(row) if row else None + + def get_documents_by_source(self, source: str, status: Optional[str] = None, + parse_status: Optional[str] = None) -> list[Document]: + query = "SELECT * FROM documents WHERE source = ?" + params: list = [source] + if status: + query += " AND status = ?" + params.append(status) + if parse_status: + query += " AND parse_status = ?" + params.append(parse_status) + cursor = self.conn.execute(query, params) + return [Document.from_row(row) for row in cursor.fetchall()] + + def get_documents_by_date_range(self, source: Optional[str] = None, + since: Optional[str] = None, + until: Optional[str] = None) -> list[Document]: + query = "SELECT * FROM documents WHERE 1=1" + params: list = [] + if source: + query += " AND source = ?" + params.append(source) + if since: + query += " AND date >= ?" + params.append(since) + if until: + query += " AND date <= ?" + params.append(until) + cursor = self.conn.execute(query, params) + return [Document.from_row(row) for row in cursor.fetchall()] + + def get_unextracted_documents(self, source: Optional[str] = None, + extractor: Optional[str] = None, + since_last_run: bool = False) -> list[Document]: + """Get documents that need extraction.""" + if extractor: + query = """ + SELECT d.* FROM documents d + LEFT JOIN extraction_status es + ON d.id = es.doc_id AND es.extractor = ? + WHERE d.parse_status = 'PARSED' + AND (es.status IS NULL OR es.status = 'PENDING') + """ + params: list = [extractor] + else: + query = "SELECT * FROM documents WHERE parse_status = 'PARSED'" + params = [] + + if source: + query += " AND d.source = ?" if extractor else " AND source = ?" + params.append(source) + if since_last_run: + query += " AND d.last_extracted_at IS NULL" if extractor else " AND last_extracted_at IS NULL" + + cursor = self.conn.execute(query, params) + return [Document.from_row(row) for row in cursor.fetchall()] + + def update_document_status(self, doc_id: str, status: Optional[str] = None, + parse_status: Optional[str] = None): + updates = [] + params: list = [] + if status is not None: + updates.append("status = ?") + params.append(status) + if parse_status is not None: + updates.append("parse_status = ?") + params.append(parse_status) + if updates: + params.append(doc_id) + self.conn.execute( + f"UPDATE documents SET {', '.join(updates)} WHERE id = ?", + params + ) + self.conn.commit() + + def update_document_extraction_time(self, doc_id: str): + self.conn.execute( + "UPDATE documents SET last_extracted_at = ? WHERE id = ?", + (datetime.utcnow().isoformat(), doc_id) + ) + self.conn.commit() + + # --- Extracted data operations --- + + def insert_extracted_tool(self, tool: ExtractedTool) -> bool: + """Insert tool if not duplicate. Returns True if inserted.""" + if not tool.dedup_key: + tool.compute_dedup_key() + d = tool.to_dict() + try: + cols = ', '.join(k for k in d if k != 'id') + placeholders = ', '.join('?' for k in d if k != 'id') + self.conn.execute( + f"INSERT INTO extracted_tools ({cols}) VALUES ({placeholders})", + [v for k, v in d.items() if k != 'id'] + ) + self.conn.commit() + return True + except sqlite3.IntegrityError: + # Duplicate dedup_key — merge source_doc_ids + cursor = self.conn.execute( + "SELECT source_doc_ids FROM extracted_tools WHERE dedup_key = ?", + (tool.dedup_key,) + ) + row = cursor.fetchone() + if row: + import json + existing_ids = json.loads(row['source_doc_ids']) + merged = list(set(existing_ids + tool.source_doc_ids)) + self.conn.execute( + "UPDATE extracted_tools SET source_doc_ids = ? WHERE dedup_key = ?", + (json.dumps(merged), tool.dedup_key) + ) + self.conn.commit() + return False + + def insert_ttp(self, ttp: TTP) -> bool: + if not ttp.dedup_key: + ttp.compute_dedup_key() + d = ttp.to_dict() + try: + cols = ', '.join(k for k in d if k != 'id') + placeholders = ', '.join('?' for k in d if k != 'id') + self.conn.execute( + f"INSERT INTO ttps ({cols}) VALUES ({placeholders})", + [v for k, v in d.items() if k != 'id'] + ) + self.conn.commit() + return True + except sqlite3.IntegrityError: + import json + cursor = self.conn.execute( + "SELECT source_doc_ids FROM ttps WHERE dedup_key = ?", + (ttp.dedup_key,) + ) + row = cursor.fetchone() + if row: + existing_ids = json.loads(row['source_doc_ids']) + merged = list(set(existing_ids + ttp.source_doc_ids)) + self.conn.execute( + "UPDATE ttps SET source_doc_ids = ? WHERE dedup_key = ?", + (json.dumps(merged), ttp.dedup_key) + ) + self.conn.commit() + return False + + def insert_tradecraft(self, tc: Tradecraft): + d = tc.to_dict() + cols = ', '.join(k for k in d if k != 'id') + placeholders = ', '.join('?' for k in d if k != 'id') + self.conn.execute( + f"INSERT INTO tradecraft ({cols}) VALUES ({placeholders})", + [v for k, v in d.items() if k != 'id'] + ) + self.conn.commit() + + def insert_entity(self, entity: Entity): + d = entity.to_dict() + cols = ', '.join(k for k in d if k != 'id') + placeholders = ', '.join('?' for k in d if k != 'id') + self.conn.execute( + f"INSERT INTO entities ({cols}) VALUES ({placeholders})", + [v for k, v in d.items() if k != 'id'] + ) + self.conn.commit() + + def insert_infrastructure(self, infra: Infrastructure): + d = infra.to_dict() + cols = ', '.join(k for k in d if k != 'id') + placeholders = ', '.join('?' for k in d if k != 'id') + self.conn.execute( + f"INSERT INTO infrastructure ({cols}) VALUES ({placeholders})", + [v for k, v in d.items() if k != 'id'] + ) + self.conn.commit() + + # --- Extraction status --- + + def get_extraction_status(self, doc_id: str, extractor: str) -> Optional[ExtractionStatus]: + cursor = self.conn.execute( + "SELECT * FROM extraction_status WHERE doc_id = ? AND extractor = ?", + (doc_id, extractor) + ) + row = cursor.fetchone() + return ExtractionStatus.from_row(row) if row else None + + def set_extraction_status(self, doc_id: str, extractor: str, + status: str, error: Optional[str] = None, + token_count: int = 0): + self.conn.execute(""" + INSERT INTO extraction_status (doc_id, extractor, status, error, last_run, token_count) + VALUES (?, ?, ?, ?, ?, ?) + ON CONFLICT(doc_id, extractor) + DO UPDATE SET status = ?, error = ?, last_run = ?, token_count = ? + """, ( + doc_id, extractor, status, error, datetime.utcnow().isoformat(), token_count, + status, error, datetime.utcnow().isoformat(), token_count + )) + self.conn.commit() + + # --- Token usage --- + + def log_token_usage(self, usage: TokenUsage): + d = usage.to_dict() + cols = ', '.join(k for k in d if k != 'id') + placeholders = ', '.join('?' for k in d if k != 'id') + self.conn.execute( + f"INSERT INTO token_usage ({cols}) VALUES ({placeholders})", + [v for k, v in d.items() if k != 'id'] + ) + self.conn.commit() + + def get_total_tokens(self, source: Optional[str] = None) -> dict: + if source: + cursor = self.conn.execute( + "SELECT SUM(input_tokens) as input, SUM(output_tokens) as output FROM token_usage WHERE source = ?", + (source,) + ) + else: + cursor = self.conn.execute( + "SELECT SUM(input_tokens) as input, SUM(output_tokens) as output FROM token_usage" + ) + row = cursor.fetchone() + return { + 'input_tokens': row['input'] or 0, + 'output_tokens': row['output'] or 0, + 'total_tokens': (row['input'] or 0) + (row['output'] or 0) + } + + # --- FTS5 search --- + + def search_documents(self, query: str, source: Optional[str] = None, + limit: int = 50) -> list[Document]: + """Full-text search with sanitized query.""" + sanitized = self._sanitize_fts_query(query) + if not sanitized: + return [] + + if source: + sql = """ + SELECT d.* FROM documents d + JOIN documents_fts fts ON d.rowid = fts.rowid + WHERE documents_fts MATCH ? AND d.source = ? + ORDER BY rank + LIMIT ? + """ + params = [sanitized, source, limit] + else: + sql = """ + SELECT d.* FROM documents d + JOIN documents_fts fts ON d.rowid = fts.rowid + WHERE documents_fts MATCH ? + ORDER BY rank + LIMIT ? + """ + params = [sanitized, limit] + + try: + cursor = self.conn.execute(sql, params) + return [Document.from_row(row) for row in cursor.fetchall()] + except sqlite3.OperationalError as e: + logger.warning("FTS query failed: %s (query: %s)", e, sanitized) + return [] + + def _sanitize_fts_query(self, query: str) -> str: + """Sanitize FTS5 query to prevent injection and complexity issues.""" + if len(query) > MAX_FTS_QUERY_LENGTH: + query = query[:MAX_FTS_QUERY_LENGTH] + # Escape special FTS5 characters + sanitized = FTS5_SPECIAL.sub(' ', query) + # Remove excess whitespace + sanitized = ' '.join(sanitized.split()) + return sanitized.strip() + + # --- Status / stats --- + + def get_status_counts(self, source: Optional[str] = None) -> dict: + """Get document counts by status.""" + base = "SELECT status, parse_status, COUNT(*) as cnt FROM documents" + params: list = [] + if source: + base += " WHERE source = ?" + params.append(source) + base += " GROUP BY status, parse_status" + cursor = self.conn.execute(base, params) + result = { + 'collection': {}, + 'parsing': {}, + 'total': 0 + } + for row in cursor.fetchall(): + status = row['status'] + parse_status = row['parse_status'] + count = row['cnt'] + result['collection'][status] = result['collection'].get(status, 0) + count + result['parsing'][parse_status] = result['parsing'].get(parse_status, 0) + count + result['total'] += count + return result + + def get_extraction_counts(self, source: Optional[str] = None) -> dict: + """Get extraction counts by extractor and status.""" + if source: + cursor = self.conn.execute(""" + SELECT es.extractor, es.status, COUNT(*) as cnt + FROM extraction_status es + JOIN documents d ON d.id = es.doc_id + WHERE d.source = ? + GROUP BY es.extractor, es.status + """, (source,)) + else: + cursor = self.conn.execute(""" + SELECT extractor, status, COUNT(*) as cnt + FROM extraction_status + GROUP BY extractor, status + """) + result = {} + for row in cursor.fetchall(): + ext = row['extractor'] + if ext not in result: + result[ext] = {} + result[ext][row['status']] = row['cnt'] + return result + + def get_all_tools(self) -> list[ExtractedTool]: + cursor = self.conn.execute("SELECT * FROM extracted_tools ORDER BY name") + return [ExtractedTool.from_row(row) for row in cursor.fetchall()] + + def get_all_ttps(self) -> list[TTP]: + cursor = self.conn.execute("SELECT * FROM ttps ORDER BY category, technique") + return [TTP.from_row(row) for row in cursor.fetchall()] + + def get_all_tradecraft(self) -> list[Tradecraft]: + cursor = self.conn.execute("SELECT * FROM tradecraft ORDER BY domain, method") + return [Tradecraft.from_row(row) for row in cursor.fetchall()] + + def get_all_entities(self) -> list[Entity]: + cursor = self.conn.execute("SELECT * FROM entities ORDER BY type, name") + return [Entity.from_row(row) for row in cursor.fetchall()] + + def get_all_infrastructure(self) -> list[Infrastructure]: + cursor = self.conn.execute("SELECT * FROM infrastructure ORDER BY indicator_type, indicator") + return [Infrastructure.from_row(row) for row in cursor.fetchall()] + + # --- Verify integrity --- + + def get_all_document_hashes(self) -> list[tuple[str, str, str]]: + """Return (id, raw_path, content_hash) for all cached documents.""" + cursor = self.conn.execute( + "SELECT id, raw_path, content_hash FROM documents WHERE status = 'CACHED' OR parse_status = 'PARSED'" + ) + return [(row['id'], row['raw_path'], row['content_hash']) for row in cursor.fetchall()] diff --git a/storage/jsonl_store.py b/storage/jsonl_store.py new file mode 100644 index 0000000..ea63d3d --- /dev/null +++ b/storage/jsonl_store.py @@ -0,0 +1,128 @@ +"""JSONL export/import for Mosaic extracted data.""" +import json +import hashlib +import logging +from pathlib import Path +from datetime import datetime +from typing import Optional + +from .db import MosaicDB + +logger = logging.getLogger(__name__) + +CATEGORIES = { + 'tools': ('extracted_tools', 'get_all_tools'), + 'ttps': ('ttps', 'get_all_ttps'), + 'tradecraft': ('tradecraft', 'get_all_tradecraft'), + 'entities': ('entities', 'get_all_entities'), + 'infrastructure': ('infrastructure', 'get_all_infrastructure'), +} + + +class JSONLStore: + """Export and import extracted data as JSONL files with checksums.""" + + def __init__(self, db: MosaicDB, output_dir: str = "output/extracted"): + self.db = db + self.output_dir = Path(output_dir) + self.output_dir.mkdir(parents=True, exist_ok=True) + + def export(self, category: Optional[str] = None, source: Optional[str] = None) -> list[str]: + """Export data to JSONL files. Returns list of created file paths.""" + categories = [category] if category else list(CATEGORIES.keys()) + created = [] + + for cat in categories: + if cat not in CATEGORIES: + logger.warning("Unknown category: %s", cat) + continue + + table_name, getter_name = CATEGORIES[cat] + getter = getattr(self.db, getter_name) + items = getter() + + if not items: + logger.info("No data for category: %s", cat) + continue + + timestamp = datetime.utcnow().strftime("%Y%m%d_%H%M%S") + filename = f"{cat}_{timestamp}.jsonl" + filepath = self.output_dir / filename + + sha256 = hashlib.sha256() + count = 0 + with open(filepath, 'w') as f: + for item in items: + line = json.dumps(item.to_dict(), ensure_ascii=False) + '\n' + f.write(line) + sha256.update(line.encode()) + count += 1 + + # Write checksum file + checksum_path = filepath.with_suffix('.jsonl.sha256') + checksum_path.write_text(f"{sha256.hexdigest()} {filename}\n") + + logger.info("Exported %d %s to %s", count, cat, filepath) + created.append(str(filepath)) + + return created + + def import_jsonl(self, filepath: str, category: str) -> int: + """Import data from a JSONL file. Returns count of imported records.""" + from .models import ExtractedTool, TTP, Tradecraft, Entity, Infrastructure + + model_map = { + 'tools': ExtractedTool, + 'ttps': TTP, + 'tradecraft': Tradecraft, + 'entities': Entity, + 'infrastructure': Infrastructure, + } + insert_map = { + 'tools': self.db.insert_extracted_tool, + 'ttps': self.db.insert_ttp, + 'tradecraft': self.db.insert_tradecraft, + 'entities': self.db.insert_entity, + 'infrastructure': self.db.insert_infrastructure, + } + + if category not in model_map: + raise ValueError(f"Unknown category: {category}") + + model_cls = model_map[category] + inserter = insert_map[category] + count = 0 + + with open(filepath, 'r') as f: + for line in f: + line = line.strip() + if not line: + continue + data = json.loads(line) + # Deserialize JSON list fields + obj = model_cls.from_row(data) + inserter(obj) + count += 1 + + logger.info("Imported %d %s from %s", count, category, filepath) + return count + + def verify_checksum(self, filepath: str) -> bool: + """Verify a JSONL file against its checksum.""" + checksum_path = Path(filepath).with_suffix('.jsonl.sha256') + if not checksum_path.exists(): + logger.warning("No checksum file found for %s", filepath) + return False + + expected = checksum_path.read_text().split()[0] + sha256 = hashlib.sha256() + with open(filepath, 'rb') as f: + for chunk in iter(lambda: f.read(8192), b''): + sha256.update(chunk) + + actual = sha256.hexdigest() + if actual != expected: + logger.error("Checksum mismatch for %s: expected %s, got %s", + filepath, expected, actual) + return False + return True diff --git a/storage/migrations.py b/storage/migrations.py new file mode 100644 index 0000000..afe6af0 --- /dev/null +++ b/storage/migrations.py @@ -0,0 +1,191 @@ +"""Schema versioning and migration runner for Mosaic database.""" +import sqlite3 +import logging + +logger = logging.getLogger(__name__) + +MIGRATIONS = [ + # Version 1: Initial schema + { + "version": 1, + "description": "Initial schema", + "sql": [ + """CREATE TABLE IF NOT EXISTS documents ( + id TEXT PRIMARY KEY, + source TEXT NOT NULL DEFAULT '', + source_url TEXT NOT NULL DEFAULT '', + doc_type TEXT NOT NULL DEFAULT '', + title TEXT NOT NULL DEFAULT '', + date TEXT, + classification TEXT, + raw_path TEXT NOT NULL DEFAULT '', + text TEXT NOT NULL DEFAULT '', + metadata TEXT NOT NULL DEFAULT '{}', + content_hash TEXT NOT NULL DEFAULT '', + fetch_date TEXT NOT NULL DEFAULT '', + etag TEXT, + last_modified TEXT, + char_count INTEGER NOT NULL DEFAULT 0, + status TEXT NOT NULL DEFAULT 'PENDING', + parse_status TEXT NOT NULL DEFAULT 'PENDING', + last_extracted_at TEXT + )""", + """CREATE INDEX IF NOT EXISTS idx_documents_source ON documents(source)""", + """CREATE INDEX IF NOT EXISTS idx_documents_status ON documents(status)""", + """CREATE INDEX IF NOT EXISTS idx_documents_parse_status ON documents(parse_status)""", + """CREATE INDEX IF NOT EXISTS idx_documents_date ON documents(date)""", + """CREATE INDEX IF NOT EXISTS idx_documents_source_url ON documents(source_url)""", + """CREATE VIRTUAL TABLE IF NOT EXISTS documents_fts USING fts5( + title, text, source, + content=documents, + content_rowid=rowid + )""", + # Triggers to keep FTS in sync + """CREATE TRIGGER IF NOT EXISTS documents_ai AFTER INSERT ON documents BEGIN + INSERT INTO documents_fts(rowid, title, text, source) + VALUES (new.rowid, new.title, new.text, new.source); + END""", + """CREATE TRIGGER IF NOT EXISTS documents_ad AFTER DELETE ON documents BEGIN + INSERT INTO documents_fts(documents_fts, rowid, title, text, source) + VALUES ('delete', old.rowid, old.title, old.text, old.source); + END""", + """CREATE TRIGGER IF NOT EXISTS documents_au AFTER UPDATE ON documents BEGIN + INSERT INTO documents_fts(documents_fts, rowid, title, text, source) + VALUES ('delete', old.rowid, old.title, old.text, old.source); + INSERT INTO documents_fts(rowid, title, text, source) + VALUES (new.rowid, new.title, new.text, new.source); + END""", + """CREATE TABLE IF NOT EXISTS extracted_tools ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + name TEXT NOT NULL DEFAULT '', + aliases TEXT NOT NULL DEFAULT '[]', + description TEXT NOT NULL DEFAULT '', + capability TEXT NOT NULL DEFAULT '', + target_platforms TEXT NOT NULL DEFAULT '[]', + target_software TEXT NOT NULL DEFAULT '[]', + cves TEXT NOT NULL DEFAULT '[]', + source_doc_ids TEXT NOT NULL DEFAULT '[]', + source_context TEXT NOT NULL DEFAULT '', + mitre_techniques TEXT NOT NULL DEFAULT '[]', + dedup_key TEXT NOT NULL DEFAULT '' + )""", + """CREATE UNIQUE INDEX IF NOT EXISTS idx_tools_dedup ON extracted_tools(dedup_key)""", + """CREATE TABLE IF NOT EXISTS ttps ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + technique TEXT NOT NULL DEFAULT '', + category TEXT NOT NULL DEFAULT '', + description TEXT NOT NULL DEFAULT '', + actor TEXT, + mitre_id TEXT, + source_doc_ids TEXT NOT NULL DEFAULT '[]', + source_context TEXT NOT NULL DEFAULT '', + dedup_key TEXT NOT NULL DEFAULT '' + )""", + """CREATE UNIQUE INDEX IF NOT EXISTS idx_ttps_dedup ON ttps(dedup_key)""", + """CREATE TABLE IF NOT EXISTS tradecraft ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + method TEXT NOT NULL DEFAULT '', + domain TEXT NOT NULL DEFAULT '', + description TEXT NOT NULL DEFAULT '', + operational_notes TEXT NOT NULL DEFAULT '', + countermeasures TEXT NOT NULL DEFAULT '', + source_doc_ids TEXT NOT NULL DEFAULT '[]', + source_context TEXT NOT NULL DEFAULT '' + )""", + """CREATE TABLE IF NOT EXISTS entities ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + name TEXT NOT NULL DEFAULT '', + type TEXT NOT NULL DEFAULT '', + description TEXT NOT NULL DEFAULT '', + relationships TEXT NOT NULL DEFAULT '[]', + source_doc_ids TEXT NOT NULL DEFAULT '[]', + source_context TEXT NOT NULL DEFAULT '' + )""", + """CREATE TABLE IF NOT EXISTS infrastructure ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + indicator TEXT NOT NULL DEFAULT '', + indicator_type TEXT NOT NULL DEFAULT '', + context TEXT NOT NULL DEFAULT '', + source_doc_ids TEXT NOT NULL DEFAULT '[]', + source_context TEXT NOT NULL DEFAULT '', + collection TEXT NOT NULL DEFAULT '', + doc_title TEXT NOT NULL DEFAULT '', + doc_date TEXT + )""", + """CREATE INDEX IF NOT EXISTS idx_infra_indicator ON infrastructure(indicator)""", + """CREATE INDEX IF NOT EXISTS idx_infra_type ON infrastructure(indicator_type)""", + """CREATE TABLE IF NOT EXISTS extraction_status ( + doc_id TEXT NOT NULL DEFAULT '', + extractor TEXT NOT NULL DEFAULT '', + status TEXT NOT NULL DEFAULT 'PENDING', + error TEXT, + last_run TEXT, + token_count INTEGER NOT NULL DEFAULT 0, + PRIMARY KEY (doc_id, extractor) + )""", + """CREATE TABLE IF NOT EXISTS token_usage ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + timestamp TEXT NOT NULL DEFAULT '', + source TEXT NOT NULL DEFAULT '', + doc_id TEXT NOT NULL DEFAULT '', + extractor TEXT NOT NULL DEFAULT '', + input_tokens INTEGER NOT NULL DEFAULT 0, + output_tokens INTEGER NOT NULL DEFAULT 0, + model TEXT NOT NULL DEFAULT '' + )""", + ] + }, +] + + +class MigrationRunner: + """Run schema migrations on a SQLite database.""" + + def __init__(self, conn: sqlite3.Connection): + self.conn = conn + self._ensure_version_table() + + def _ensure_version_table(self): + self.conn.execute(""" + CREATE TABLE IF NOT EXISTS schema_version ( + version INTEGER PRIMARY KEY, + description TEXT NOT NULL DEFAULT '', + applied_at TEXT NOT NULL DEFAULT (datetime('now')) + ) + """) + self.conn.commit() + + def get_current_version(self) -> int: + cursor = self.conn.execute( + "SELECT MAX(version) FROM schema_version" + ) + row = cursor.fetchone() + return row[0] if row[0] is not None else 0 + + def run(self): + current = self.get_current_version() + pending = [m for m in MIGRATIONS if m['version'] > current] + + if not pending: + logger.debug("Database schema is up to date (version %d)", current) + return + + for migration in sorted(pending, key=lambda m: m['version']): + logger.info( + "Applying migration %d: %s", + migration['version'], + migration['description'] + ) + try: + for sql in migration['sql']: + self.conn.execute(sql) + self.conn.execute( + "INSERT INTO schema_version (version, description) VALUES (?, ?)", + (migration['version'], migration['description']) + ) + self.conn.commit() + logger.info("Migration %d applied successfully", migration['version']) + except Exception as e: + self.conn.rollback() + logger.error("Migration %d failed: %s", migration['version'], e) + raise diff --git a/storage/models.py b/storage/models.py new file mode 100644 index 0000000..662ef4c --- /dev/null +++ b/storage/models.py @@ -0,0 +1,227 @@ +"""Data models for Mosaic intelligence extraction.""" +from dataclasses import dataclass, field, asdict +from typing import Optional +import hashlib +import json + + +@dataclass +class Document: + id: str = "" # SHA256 of content + source: str = "" # Source profile name + source_url: str = "" + doc_type: str = "" # html, pdf, email, cable, text + title: str = "" + date: Optional[str] = None + classification: Optional[str] = None + raw_path: str = "" # Hash-sharded path to raw file + text: str = "" # Extracted plain text + metadata: dict = field(default_factory=dict) + content_hash: str = "" # SHA256 verified on every access + fetch_date: str = "" + etag: Optional[str] = None + last_modified: Optional[str] = None + char_count: int = 0 + status: str = "PENDING" # PENDING|DOWNLOADING|CACHED|DOWNLOAD_FAILED|IN_PROGRESS|PARSED|PARSE_FAILED|NEEDS_OCR + parse_status: str = "PENDING" # PENDING|IN_PROGRESS|PARSED|PARSE_FAILED|NEEDS_OCR + last_extracted_at: Optional[str] = None + + def to_dict(self) -> dict: + d = asdict(self) + d['metadata'] = json.dumps(d['metadata']) + return d + + @classmethod + def from_row(cls, row: dict) -> 'Document': + row = dict(row) + if isinstance(row.get('metadata'), str): + try: + row['metadata'] = json.loads(row['metadata']) + except (json.JSONDecodeError, TypeError): + row['metadata'] = {} + return cls(**{k: v for k, v in row.items() if k in cls.__dataclass_fields__}) + + +@dataclass +class ExtractedTool: + name: str = "" + aliases: list = field(default_factory=list) + description: str = "" + capability: str = "" # exploit, implant, framework, utility + target_platforms: list = field(default_factory=list) + target_software: list = field(default_factory=list) + cves: list = field(default_factory=list) + source_doc_ids: list = field(default_factory=list) + source_context: str = "" + mitre_techniques: list = field(default_factory=list) + dedup_key: str = "" + + def compute_dedup_key(self) -> str: + raw = f"{self.name.lower().strip()}:{self.capability.lower().strip()}" + self.dedup_key = hashlib.sha256(raw.encode()).hexdigest()[:16] + return self.dedup_key + + def to_dict(self) -> dict: + d = asdict(self) + for k in ('aliases', 'target_platforms', 'target_software', 'cves', + 'source_doc_ids', 'mitre_techniques'): + d[k] = json.dumps(d[k]) + return d + + @classmethod + def from_row(cls, row: dict) -> 'ExtractedTool': + row = dict(row) + for k in ('aliases', 'target_platforms', 'target_software', 'cves', + 'source_doc_ids', 'mitre_techniques'): + if isinstance(row.get(k), str): + try: + row[k] = json.loads(row[k]) + except (json.JSONDecodeError, TypeError): + row[k] = [] + return cls(**{k: v for k, v in row.items() if k in cls.__dataclass_fields__}) + + +@dataclass +class TTP: + technique: str = "" + category: str = "" # initial_access, persistence, exfil, etc. + description: str = "" + actor: Optional[str] = None + mitre_id: Optional[str] = None + source_doc_ids: list = field(default_factory=list) + source_context: str = "" + dedup_key: str = "" + + def compute_dedup_key(self) -> str: + raw = f"{self.technique.lower().strip()}:{self.category.lower().strip()}" + self.dedup_key = hashlib.sha256(raw.encode()).hexdigest()[:16] + return self.dedup_key + + def to_dict(self) -> dict: + d = asdict(self) + d['source_doc_ids'] = json.dumps(d['source_doc_ids']) + return d + + @classmethod + def from_row(cls, row: dict) -> 'TTP': + row = dict(row) + if isinstance(row.get('source_doc_ids'), str): + try: + row['source_doc_ids'] = json.loads(row['source_doc_ids']) + except (json.JSONDecodeError, TypeError): + row['source_doc_ids'] = [] + return cls(**{k: v for k, v in row.items() if k in cls.__dataclass_fields__}) + + +@dataclass +class Tradecraft: + method: str = "" + domain: str = "" # humint, sigint, cyber, surveillance + description: str = "" + operational_notes: str = "" + countermeasures: str = "" + source_doc_ids: list = field(default_factory=list) + source_context: str = "" + + def to_dict(self) -> dict: + d = asdict(self) + d['source_doc_ids'] = json.dumps(d['source_doc_ids']) + return d + + @classmethod + def from_row(cls, row: dict) -> 'Tradecraft': + row = dict(row) + if isinstance(row.get('source_doc_ids'), str): + try: + row['source_doc_ids'] = json.loads(row['source_doc_ids']) + except (json.JSONDecodeError, TypeError): + row['source_doc_ids'] = [] + return cls(**{k: v for k, v in row.items() if k in cls.__dataclass_fields__}) + + +@dataclass +class Entity: + name: str = "" + type: str = "" # person, org, program, codename, unit + description: str = "" + relationships: list = field(default_factory=list) + source_doc_ids: list = field(default_factory=list) + source_context: str = "" + + def to_dict(self) -> dict: + d = asdict(self) + d['relationships'] = json.dumps(d['relationships']) + d['source_doc_ids'] = json.dumps(d['source_doc_ids']) + return d + + @classmethod + def from_row(cls, row: dict) -> 'Entity': + row = dict(row) + for k in ('relationships', 'source_doc_ids'): + if isinstance(row.get(k), str): + try: + row[k] = json.loads(row[k]) + except (json.JSONDecodeError, TypeError): + row[k] = [] + return cls(**{k: v for k, v in row.items() if k in cls.__dataclass_fields__}) + + +@dataclass +class Infrastructure: + indicator: str = "" + indicator_type: str = "" # ipv4, ipv6, domain, url, hash + context: str = "" + source_doc_ids: list = field(default_factory=list) + source_context: str = "" + collection: str = "" + doc_title: str = "" + doc_date: Optional[str] = None + + def to_dict(self) -> dict: + d = asdict(self) + d['source_doc_ids'] = json.dumps(d['source_doc_ids']) + return d + + @classmethod + def from_row(cls, row: dict) -> 'Infrastructure': + row = dict(row) + if isinstance(row.get('source_doc_ids'), str): + try: + row['source_doc_ids'] = json.loads(row['source_doc_ids']) + except (json.JSONDecodeError, TypeError): + row['source_doc_ids'] = [] + return cls(**{k: v for k, v in row.items() if k in cls.__dataclass_fields__}) + + +@dataclass +class ExtractionStatus: + """Tracks per-document, per-extractor status.""" + doc_id: str = "" + extractor: str = "" # tool, ttp, tradecraft, surveillance, infrastructure, entity + status: str = "PENDING" # PENDING|IN_PROGRESS|DONE|FAILED + error: Optional[str] = None + last_run: Optional[str] = None + token_count: int = 0 + + def to_dict(self) -> dict: + return asdict(self) + + @classmethod + def from_row(cls, row: dict) -> 'ExtractionStatus': + row = dict(row) + return cls(**{k: v for k, v in row.items() if k in cls.__dataclass_fields__}) + + +@dataclass +class TokenUsage: + """Track API token usage per call.""" + timestamp: str = "" + source: str = "" + doc_id: str = "" + extractor: str = "" + input_tokens: int = 0 + output_tokens: int = 0 + model: str = "" + + def to_dict(self) -> dict: + return asdict(self) diff --git a/tests/.gitkeep b/tests/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/utils/__init__.py b/utils/__init__.py new file mode 100644 index 0000000..d1627ec --- /dev/null +++ b/utils/__init__.py @@ -0,0 +1,5 @@ +from .http_limiter import HTTPRateLimiter +from .api_limiter import APIRateLimiter +from .cache import DownloadCache +from .sanitize import sanitize_filename, sanitize_path, validate_content_type +from .text_utils import chunk_text, clean_text, estimate_tokens diff --git a/utils/api_limiter.py b/utils/api_limiter.py new file mode 100644 index 0000000..8106923 --- /dev/null +++ b/utils/api_limiter.py @@ -0,0 +1,92 @@ +"""Exponential backoff rate limiter with circuit breaker for Anthropic API.""" +import time +import random +import threading +import logging +from typing import Optional + +logger = logging.getLogger(__name__) + + +class CircuitBreakerOpen(Exception): + """Raised when circuit breaker is open (too many consecutive failures).""" + pass + + +class APIRateLimiter: + """Exponential backoff with jitter and circuit breaker for API calls.""" + + def __init__(self, base_delay: float = 1.0, max_delay: float = 60.0, + circuit_threshold: int = 5): + """ + Args: + base_delay: Initial backoff delay in seconds + max_delay: Maximum backoff delay in seconds + circuit_threshold: Consecutive failures before circuit opens + """ + self.base_delay = base_delay + self.max_delay = max_delay + self.circuit_threshold = circuit_threshold + self._consecutive_failures = 0 + self._current_delay = 0.0 + self._lock = threading.Lock() + + @property + def is_circuit_open(self) -> bool: + return self._consecutive_failures >= self.circuit_threshold + + def record_success(self): + """Record a successful API call — reset backoff.""" + with self._lock: + self._consecutive_failures = 0 + self._current_delay = 0.0 + + def record_failure(self, error: Optional[str] = None): + """Record a failed API call — increase backoff.""" + with self._lock: + self._consecutive_failures += 1 + logger.warning( + "API failure #%d/%d: %s", + self._consecutive_failures, + self.circuit_threshold, + error or "unknown" + ) + if self.is_circuit_open: + logger.error( + "Circuit breaker OPEN: %d consecutive failures. " + "Halting API calls until manual reset.", + self._consecutive_failures + ) + + def wait_if_needed(self): + """Wait with exponential backoff if there have been failures. Raises CircuitBreakerOpen.""" + with self._lock: + if self.is_circuit_open: + raise CircuitBreakerOpen( + f"Circuit breaker open after {self._consecutive_failures} consecutive failures. " + f"Call reset() to retry." + ) + + if self._consecutive_failures == 0: + return + + delay = min( + self.base_delay * (2 ** (self._consecutive_failures - 1)), + self.max_delay + ) + # Add jitter (±25%) + jitter = delay * random.uniform(-0.25, 0.25) + actual_delay = max(0, delay + jitter) + + logger.info("API backoff: waiting %.1fs (attempt %d)", + actual_delay, self._consecutive_failures + 1) + + # Sleep outside the lock + time.sleep(actual_delay) + + def reset(self): + """Manually reset the circuit breaker.""" + with self._lock: + self._consecutive_failures = 0 + self._current_delay = 0.0 + logger.info("Circuit breaker reset") diff --git a/utils/cache.py b/utils/cache.py new file mode 100644 index 0000000..9f8686f --- /dev/null +++ b/utils/cache.py @@ -0,0 +1,147 @@ +"""Atomic download cache with hash verification for Mosaic.""" +import os +import hashlib +import shutil +import logging +from pathlib import Path +from typing import Optional + +logger = logging.getLogger(__name__) + + +class DownloadCache: + """Manages atomic downloads with hash-sharded storage and integrity verification.""" + + def __init__(self, base_dir: str = "output/raw"): + self.base_dir = Path(base_dir) + self.base_dir.mkdir(parents=True, exist_ok=True) + + def get_shard_path(self, content_hash: str, collection: str, ext: str = "") -> Path: + """Get hash-sharded storage path: raw/{collection}/{ab}/{cd}/{hash}.{ext}""" + if len(content_hash) < 4: + raise ValueError(f"Invalid hash: {content_hash}") + ab = content_hash[:2] + cd = content_hash[2:4] + filename = content_hash + (f".{ext}" if ext else "") + return self.base_dir / collection / ab / cd / filename + + def store_atomic(self, tmp_path: str, content_hash: str, collection: str, + ext: str = "", expected_size: Optional[int] = None) -> str: + """Atomically move a temp file to its cache location. + + Args: + tmp_path: Path to the temporary download file + content_hash: SHA256 hash of the content + collection: Source collection name + ext: File extension + expected_size: Expected file size (from Content-Length) + + Returns: + Final path of the cached file + + Raises: + ValueError: If hash or size verification fails + """ + tmp = Path(tmp_path) + if not tmp.exists(): + raise FileNotFoundError(f"Temp file not found: {tmp_path}") + + # Verify size if expected + if expected_size is not None: + actual_size = tmp.stat().st_size + if actual_size != expected_size: + raise ValueError( + f"Size mismatch: expected {expected_size}, got {actual_size}" + ) + + # Verify hash + actual_hash = self.compute_hash(tmp_path) + if actual_hash != content_hash: + raise ValueError( + f"Hash mismatch: expected {content_hash}, got {actual_hash}" + ) + + # Atomic move to final location + final_path = self.get_shard_path(content_hash, collection, ext) + final_path.parent.mkdir(parents=True, exist_ok=True) + + shutil.move(str(tmp), str(final_path)) + logger.debug("Cached: %s → %s", tmp_path, final_path) + return str(final_path) + + def is_cached(self, content_hash: str, collection: str, ext: str = "") -> bool: + """Check if a file exists in cache with valid hash.""" + path = self.get_shard_path(content_hash, collection, ext) + if not path.exists(): + return False + actual_hash = self.compute_hash(str(path)) + if actual_hash != content_hash: + logger.warning("Cache corruption detected: %s (expected %s, got %s)", + path, content_hash, actual_hash) + return False + return True + + def get_cached_path(self, content_hash: str, collection: str, ext: str = "") -> Optional[str]: + """Get path to cached file, or None if not cached/corrupted.""" + if self.is_cached(content_hash, collection, ext): + return str(self.get_shard_path(content_hash, collection, ext)) + return None + + def verify_file(self, file_path: str, expected_hash: str) -> bool: + """Verify a file's hash matches expected.""" + actual = self.compute_hash(file_path) + return actual == expected_hash + + def verify_all(self, collection: str) -> tuple[int, int, list[str]]: + """Verify all files in a collection. + + Returns: + (verified_count, mismatch_count, list_of_mismatched_paths) + """ + collection_dir = self.base_dir / collection + if not collection_dir.exists(): + return 0, 0, [] + + verified = 0 + mismatched = 0 + bad_paths = [] + + for root, _, files in os.walk(collection_dir): + for fname in files: + if fname.startswith('.') or fname.endswith('.tmp'): + continue + fpath = os.path.join(root, fname) + # Extract expected hash from filename + expected_hash = Path(fname).stem + actual_hash = self.compute_hash(fpath) + if actual_hash == expected_hash: + verified += 1 + else: + mismatched += 1 + bad_paths.append(fpath) + + return verified, mismatched, bad_paths + + @staticmethod + def compute_hash(file_path: str) -> str: + """Compute SHA256 hash of a file.""" + sha256 = hashlib.sha256() + with open(file_path, 'rb') as f: + for chunk in iter(lambda: f.read(8192), b''): + sha256.update(chunk) + return sha256.hexdigest() + + @staticmethod + def hash_content(content: bytes) -> str: + """Compute SHA256 hash of bytes content.""" + return hashlib.sha256(content).hexdigest() + + def check_disk_space(self, required_mb: int = 100) -> bool: + """Check if enough disk space is available.""" + stat = shutil.disk_usage(str(self.base_dir)) + free_mb = stat.free // (1024 * 1024) + if free_mb < required_mb: + logger.warning("Low disk space: %d MB free, %d MB required", + free_mb, required_mb) + return False + return True diff --git a/utils/http_limiter.py b/utils/http_limiter.py new file mode 100644 index 0000000..60e0d0e --- /dev/null +++ b/utils/http_limiter.py @@ -0,0 +1,45 @@ +"""Token-bucket rate limiter for HTTP requests.""" +import time +import random +import threading +import logging + +logger = logging.getLogger(__name__) + + +class HTTPRateLimiter: + """Token-bucket rate limiter with jitter for polite crawling.""" + + def __init__(self, rate: float = 1.0, jitter: float = 0.5): + """ + Args: + rate: Maximum requests per second + jitter: Random jitter range in seconds (±jitter/2) + """ + self.rate = rate + self.jitter = jitter + self.min_interval = 1.0 / rate if rate > 0 else 0 + self._last_request = 0.0 + self._lock = threading.Lock() + + def wait(self): + """Block until the next request is allowed.""" + with self._lock: + now = time.monotonic() + elapsed = now - self._last_request + wait_time = self.min_interval - elapsed + + if self.jitter > 0: + wait_time += random.uniform(-self.jitter / 2, self.jitter / 2) + + if wait_time > 0: + logger.debug("Rate limiting: waiting %.2fs", wait_time) + time.sleep(wait_time) + + self._last_request = time.monotonic() + + def update_rate(self, rate: float): + """Update the rate limit (e.g., from profile config).""" + with self._lock: + self.rate = rate + self.min_interval = 1.0 / rate if rate > 0 else 0 diff --git a/utils/sanitize.py b/utils/sanitize.py new file mode 100644 index 0000000..8de0cc9 --- /dev/null +++ b/utils/sanitize.py @@ -0,0 +1,141 @@ +"""Input sanitization for Mosaic — filenames, paths, content types.""" +import os +import re +import logging +from pathlib import Path +from typing import Optional + +logger = logging.getLogger(__name__) + +# Allowed characters in filenames +FILENAME_SAFE = re.compile(r'[^a-zA-Z0-9._-]') +# Maximum filename length +MAX_FILENAME_LENGTH = 200 +# ANSI escape sequence pattern +ANSI_ESCAPE = re.compile(r'\x1b\[[0-9;]*[a-zA-Z]|\x1b\].*?\x07') + +# Magic bytes → MIME type mapping +MAGIC_BYTES = { + b'%PDF': 'application/pdf', + b'\x89PNG': 'image/png', + b'\xff\xd8\xff': 'image/jpeg', + b'GIF87a': 'image/gif', + b'GIF89a': 'image/gif', + b'PK\x03\x04': 'application/zip', + b' str: + """Sanitize a filename to safe characters only. + + - Strip to basename (no directory traversal) + - Replace unsafe characters with underscore + - Limit length + """ + # Strip to basename + filename = os.path.basename(filename) + # Replace unsafe characters + sanitized = FILENAME_SAFE.sub('_', filename) + # Collapse multiple underscores + sanitized = re.sub(r'_+', '_', sanitized) + # Limit length + if len(sanitized) > MAX_FILENAME_LENGTH: + name, ext = os.path.splitext(sanitized) + sanitized = name[:MAX_FILENAME_LENGTH - len(ext)] + ext + # Ensure non-empty + if not sanitized or sanitized == '_': + sanitized = 'unnamed' + return sanitized + + +def sanitize_path(path: str, base_dir: str) -> Optional[str]: + """Validate a path stays within base_dir (path traversal prevention). + + Returns the resolved path if safe, None if path escapes base_dir. + """ + base = os.path.realpath(base_dir) + resolved = os.path.realpath(os.path.join(base, path)) + if resolved.startswith(base + os.sep) or resolved == base: + return resolved + logger.warning("Path traversal blocked: %s → %s (base: %s)", path, resolved, base) + return None + + +def validate_content_type(file_path: str, expected_type: Optional[str] = None) -> str: + """Determine content type by magic bytes, not file extension. + + Args: + file_path: Path to the file + expected_type: Expected MIME type to validate against + + Returns: + Detected MIME type string + + Raises: + ValueError: If expected_type is provided and doesn't match + """ + try: + import magic + detected = magic.from_file(file_path, mime=True) + except ImportError: + # Fallback to magic bytes check + detected = _detect_by_magic_bytes(file_path) + + if expected_type and detected != expected_type: + # Allow some flexibility (e.g., text/plain for text/html) + if not _types_compatible(detected, expected_type): + raise ValueError( + f"Content type mismatch: expected {expected_type}, detected {detected}" + ) + return detected + + +def _detect_by_magic_bytes(file_path: str) -> str: + """Fallback detection using magic bytes.""" + try: + with open(file_path, 'rb') as f: + header = f.read(16) + except (IOError, OSError): + return 'application/octet-stream' + + for magic, mime in MAGIC_BYTES.items(): + if header.startswith(magic): + return mime + + # Check if it looks like text + try: + with open(file_path, 'r', encoding='utf-8') as f: + f.read(1024) + return 'text/plain' + except (UnicodeDecodeError, IOError): + return 'application/octet-stream' + + +def _types_compatible(detected: str, expected: str) -> bool: + """Check if two MIME types are compatible.""" + # Same type + if detected == expected: + return True + # text/plain is compatible with text/* + if detected == 'text/plain' and expected.startswith('text/'): + return True + # text/html variants + if detected.startswith('text/html') and expected.startswith('text/html'): + return True + return False + + +def strip_ansi(text: str) -> str: + """Strip ANSI escape sequences from text.""" + return ANSI_ESCAPE.sub('', text) + + +def sanitize_for_display(text: str, max_length: int = 0) -> str: + """Sanitize text for safe Rich display — strip ANSI, optionally truncate.""" + clean = strip_ansi(text) + if max_length > 0 and len(clean) > max_length: + clean = clean[:max_length] + "..." + return clean diff --git a/utils/text_utils.py b/utils/text_utils.py new file mode 100644 index 0000000..5454bbc --- /dev/null +++ b/utils/text_utils.py @@ -0,0 +1,253 @@ +"""Text chunking, cleaning, and token estimation for Mosaic.""" +import re +import hashlib +import logging +from typing import Optional + +logger = logging.getLogger(__name__) + +# Rough token estimation: ~4 chars per token for English text +CHARS_PER_TOKEN = 4 + +# Default chunk sizes +DEFAULT_CHUNK_TOKENS = 8000 +DEFAULT_OVERLAP_TOKENS = 500 +EMAIL_MAX_TOKENS = 4000 +PDF_MAX_TOKENS = 8000 + + +def estimate_tokens(text: str) -> int: + """Estimate token count from text length.""" + return len(text) // CHARS_PER_TOKEN + + +def clean_text(text: str) -> str: + """Clean extracted text: normalize whitespace, remove control chars.""" + if not text: + return "" + # Remove null bytes and other control characters (keep newlines, tabs) + text = re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]', '', text) + # Normalize line endings + text = text.replace('\r\n', '\n').replace('\r', '\n') + # Collapse excessive blank lines (3+ → 2) + text = re.sub(r'\n{3,}', '\n\n', text) + # Strip trailing whitespace per line + text = '\n'.join(line.rstrip() for line in text.split('\n')) + return text.strip() + + +def chunk_text(text: str, doc_type: str = "text", + max_tokens: int = DEFAULT_CHUNK_TOKENS, + overlap_tokens: int = DEFAULT_OVERLAP_TOKENS, + metadata: Optional[dict] = None) -> list[dict]: + """Chunk text using per-document-type strategy. + + Args: + text: The text to chunk + doc_type: Document type (html, pdf, email, cable, text) + max_tokens: Maximum tokens per chunk + overlap_tokens: Token overlap between chunks + metadata: Optional metadata (e.g., cable headers) + + Returns: + List of dicts with 'text', 'chunk_index', 'total_chunks', 'metadata' + """ + if not text: + return [] + + if doc_type == 'cable': + return _chunk_cable(text, max_tokens, overlap_tokens, metadata) + elif doc_type == 'pdf': + return _chunk_pdf(text, max_tokens, overlap_tokens) + elif doc_type == 'email': + return _chunk_email(text, max_tokens, overlap_tokens, metadata) + elif doc_type == 'html': + return _chunk_html(text, max_tokens, overlap_tokens) + else: + return _chunk_default(text, max_tokens, overlap_tokens) + + +def _chunk_cable(text: str, max_tokens: int, overlap_tokens: int, + metadata: Optional[dict] = None) -> list[dict]: + """Chunk cable: headers as structured context, body chunked independently.""" + headers = metadata or {} + header_text = "" + if headers: + header_lines = [] + for key in ('SUBJECT', 'ORIGIN', 'CLASSIFICATION', 'TAGS', 'DATE'): + if key in headers: + header_lines.append(f"{key}: {headers[key]}") + header_text = '\n'.join(header_lines) + + # If small enough, return as one chunk + total_tokens = estimate_tokens(header_text + '\n\n' + text) + if total_tokens <= max_tokens: + return [{'text': (header_text + '\n\n' + text).strip(), + 'chunk_index': 0, 'total_chunks': 1, 'metadata': headers}] + + # Chunk body, prepend headers to each chunk + body_chunks = _chunk_default(text, max_tokens - estimate_tokens(header_text) - 50, + overlap_tokens) + for i, chunk in enumerate(body_chunks): + chunk['text'] = header_text + '\n\n' + chunk['text'] + chunk['metadata'] = headers + chunk['chunk_index'] = i + chunk['total_chunks'] = len(body_chunks) + return body_chunks + + +def _chunk_pdf(text: str, max_tokens: int, overlap_tokens: int) -> list[dict]: + """Chunk PDF by page boundaries (indicated by form feeds or page markers).""" + # Split by form feed or common page break markers + pages = re.split(r'\f|\n-{3,}\n|\n={3,}\n', text) + pages = [p.strip() for p in pages if p.strip()] + + if not pages: + return _chunk_default(text, max_tokens, overlap_tokens) + + chunks = [] + current_text = "" + for page in pages: + combined_tokens = estimate_tokens(current_text + '\n\n' + page) + if combined_tokens > max_tokens and current_text: + chunks.append(current_text.strip()) + # Keep overlap from end of previous chunk + overlap_chars = overlap_tokens * CHARS_PER_TOKEN + current_text = current_text[-overlap_chars:] + '\n\n' + page + else: + current_text = (current_text + '\n\n' + page).strip() + + if current_text.strip(): + chunks.append(current_text.strip()) + + return [{'text': c, 'chunk_index': i, 'total_chunks': len(chunks), 'metadata': {}} + for i, c in enumerate(chunks)] + + +def _chunk_email(text: str, max_tokens: int, overlap_tokens: int, + metadata: Optional[dict] = None) -> list[dict]: + """Chunk email: headers + body as atomic unit unless body is too large.""" + headers = metadata or {} + total_tokens = estimate_tokens(text) + + if total_tokens <= EMAIL_MAX_TOKENS: + return [{'text': text, 'chunk_index': 0, 'total_chunks': 1, 'metadata': headers}] + + # Split headers from body + header_section = "" + body = text + if '\n\n' in text: + parts = text.split('\n\n', 1) + header_section = parts[0] + body = parts[1] if len(parts) > 1 else "" + + # Chunk body, prepend headers + body_chunks = _chunk_default(body, max_tokens - estimate_tokens(header_section) - 50, + overlap_tokens) + for i, chunk in enumerate(body_chunks): + chunk['text'] = header_section + '\n\n' + chunk['text'] + chunk['metadata'] = headers + chunk['chunk_index'] = i + chunk['total_chunks'] = len(body_chunks) + return body_chunks + + +def _chunk_html(text: str, max_tokens: int, overlap_tokens: int) -> list[dict]: + """Chunk HTML: section-aware chunking preserving heading hierarchy.""" + # Split on headings + sections = re.split(r'(?=\n#{1,6}\s|\n[A-Z][A-Za-z\s]{3,}\n[=-]+\n)', text) + sections = [s for s in sections if s.strip()] + + if not sections or len(sections) == 1: + return _chunk_default(text, max_tokens, overlap_tokens) + + chunks = [] + current_text = "" + for section in sections: + combined_tokens = estimate_tokens(current_text + '\n\n' + section) + if combined_tokens > max_tokens and current_text: + chunks.append(current_text.strip()) + overlap_chars = overlap_tokens * CHARS_PER_TOKEN + current_text = current_text[-overlap_chars:] + '\n\n' + section + else: + current_text = (current_text + '\n\n' + section).strip() + + if current_text.strip(): + chunks.append(current_text.strip()) + + return [{'text': c, 'chunk_index': i, 'total_chunks': len(chunks), 'metadata': {}} + for i, c in enumerate(chunks)] + + +def _chunk_default(text: str, max_tokens: int, overlap_tokens: int) -> list[dict]: + """Default chunking: paragraph-aware with token overlap.""" + total_tokens = estimate_tokens(text) + if total_tokens <= max_tokens: + return [{'text': text, 'chunk_index': 0, 'total_chunks': 1, 'metadata': {}}] + + max_chars = max_tokens * CHARS_PER_TOKEN + overlap_chars = overlap_tokens * CHARS_PER_TOKEN + + # Split into paragraphs + paragraphs = text.split('\n\n') + chunks = [] + current = "" + + for para in paragraphs: + if len(current) + len(para) + 2 > max_chars and current: + chunks.append(current.strip()) + # Keep overlap + current = current[-overlap_chars:] + '\n\n' + para + else: + current = (current + '\n\n' + para) if current else para + + if current.strip(): + chunks.append(current.strip()) + + # If still too large (single giant paragraph), force-split + final_chunks = [] + for chunk in chunks: + if len(chunk) > max_chars * 1.5: + # Force split at sentence boundaries or character limit + for sub in _force_split(chunk, max_chars, overlap_chars): + final_chunks.append(sub) + else: + final_chunks.append(chunk) + + return [{'text': c, 'chunk_index': i, 'total_chunks': len(final_chunks), 'metadata': {}} + for i, c in enumerate(final_chunks)] + + +def _force_split(text: str, max_chars: int, overlap_chars: int) -> list[str]: + """Force-split text at sentence boundaries or character limits.""" + sentences = re.split(r'(?<=[.!?])\s+', text) + chunks = [] + current = "" + + for sent in sentences: + if len(current) + len(sent) + 1 > max_chars and current: + chunks.append(current.strip()) + current = current[-overlap_chars:] + ' ' + sent + else: + current = (current + ' ' + sent) if current else sent + + if current.strip(): + chunks.append(current.strip()) + + return chunks + + +def dedup_text(texts: list[str], threshold: float = 0.95) -> list[str]: + """Remove near-duplicate texts using hash comparison. + + For exact duplicates. Near-duplicate detection would require + more sophisticated approaches (minhash, simhash, etc.) + """ + seen = set() + unique = [] + for text in texts: + h = hashlib.md5(text.encode()).hexdigest() + if h not in seen: + seen.add(h) + unique.append(text) + return unique