Initial scaffold: .gitignore, config, requirements, data seeds, directory structure

This commit is contained in:
n0mad1k
2026-03-19 08:34:32 -04:00
commit aa0f35803f
26 changed files with 1908 additions and 0 deletions
+9
View File
@@ -0,0 +1,9 @@
config/mosaic.conf
output/
*.db
*.db-wal
*.db-shm
.env
__pycache__/
.venv/
*.pyc
View File
View File
+20
View File
@@ -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: []
+21
View File
@@ -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: []
+27
View File
@@ -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: []
+23
View File
@@ -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: []
+27
View File
@@ -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/"
+27
View File
@@ -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/"
+44
View File
@@ -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"}
]
}
+22
View File
@@ -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": []
}
View File
View File
+8
View File
@@ -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
+4
View File
@@ -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
+447
View File
@@ -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()]
+128
View File
@@ -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
+191
View File
@@ -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
+227
View File
@@ -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)
View File
+5
View File
@@ -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
+92
View File
@@ -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")
+147
View File
@@ -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
+45
View File
@@ -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
+141
View File
@@ -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'<!DOCTYPE': 'text/html',
b'<html': 'text/html',
b'<HTML': 'text/html',
}
def sanitize_filename(filename: str) -> 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
+253
View File
@@ -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