n0mad1k
129 lines
4.4 KiB
Python
129 lines
4.4 KiB
Python
"""JSONL export/import for Mosaic extracted data."""
|
|
import json
|
|
import hashlib
|
|
import logging
|
|
from pathlib import Path
|
|
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.datetime.now(datetime.timezone.utc).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
|