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

2026-03-19 08:34:32 -04:00
commit aa0f35803f
26 changed files with 1908 additions and 0 deletions
@@ -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