Initial scaffold: .gitignore, config, requirements, data seeds, directory structure
This commit is contained in:
n0mad1k
@@ -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
@@ -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()]
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user