n0mad1k
448 lines
17 KiB
Python
448 lines
17 KiB
Python
"""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()]
|