Add core test suite for storage, utils, models, and parsers

2026-03-19 08:53:12 -04:00
parent ca53524118
commit 10435118a9
1 changed files with 500 additions and 0 deletions
@@ -0,0 +1,500 @@
+"""Core tests for Mosaic intelligence extraction platform."""
+import os
+import sys
+import json
+import tempfile
+import sqlite3
+import pytest
+from pathlib import Path
+
+# Add project root to path
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+
+class TestModels:
+    """Test data model serialization and deserialization."""
+
+    def test_document_to_dict(self):
+        from storage.models import Document
+        doc = Document(
+            id="abc123",
+            source="vault7",
+            source_url="https://example.com/doc",
+            doc_type="html",
+            title="Test Doc",
+            metadata={"key": "value"},
+            content_hash="abc123",
+            status="CACHED",
+        )
+        d = doc.to_dict()
+        assert d['id'] == 'abc123'
+        assert d['source'] == 'vault7'
+        assert d['metadata'] == '{"key": "value"}'  # JSON serialized
+
+    def test_document_from_row(self):
+        from storage.models import Document
+        row = {
+            'id': 'abc123',
+            'source': 'vault7',
+            'source_url': 'https://example.com',
+            'doc_type': 'html',
+            'title': 'Test',
+            'date': None,
+            'classification': None,
+            'raw_path': '/tmp/test',
+            'text': 'hello',
+            'metadata': '{"key": "value"}',
+            'content_hash': 'abc123',
+            'fetch_date': '2024-01-01',
+            'etag': None,
+            'last_modified': None,
+            'char_count': 5,
+            'status': 'CACHED',
+            'parse_status': 'PENDING',
+            'last_extracted_at': None,
+        }
+        doc = Document.from_row(row)
+        assert doc.id == 'abc123'
+        assert doc.metadata == {'key': 'value'}
+
+    def test_extracted_tool_dedup_key(self):
+        from storage.models import ExtractedTool
+        tool = ExtractedTool(name="HIVE", capability="implant")
+        key = tool.compute_dedup_key()
+        assert len(key) == 16
+        # Same name+capability should produce same key
+        tool2 = ExtractedTool(name="hive", capability="IMPLANT")
+        assert tool2.compute_dedup_key() == key
+
+    def test_ttp_dedup_key(self):
+        from storage.models import TTP
+        ttp = TTP(technique="Spearphishing", category="initial_access")
+        key = ttp.compute_dedup_key()
+        assert len(key) == 16
+
+    def test_extracted_tool_json_fields(self):
+        from storage.models import ExtractedTool
+        tool = ExtractedTool(
+            name="Test",
+            aliases=["t1", "t2"],
+            target_platforms=["windows"],
+            cves=["CVE-2024-001"],
+            source_doc_ids=["doc1", "doc2"],
+        )
+        d = tool.to_dict()
+        assert json.loads(d['aliases']) == ["t1", "t2"]
+        assert json.loads(d['source_doc_ids']) == ["doc1", "doc2"]
+
+        # Round-trip
+        restored = ExtractedTool.from_row(d)
+        assert restored.aliases == ["t1", "t2"]
+        assert restored.source_doc_ids == ["doc1", "doc2"]
+
+
+class TestDatabase:
+    """Test SQLite database operations."""
+
+    def setup_method(self):
+        self.tmp = tempfile.mktemp(suffix='.db')
+        from storage.db import MosaicDB
+        self.db = MosaicDB(self.tmp)
+
+    def teardown_method(self):
+        self.db.close()
+        for ext in ('', '-wal', '-shm'):
+            path = self.tmp + ext
+            if os.path.exists(path):
+                os.unlink(path)
+
+    def test_wal_mode(self):
+        cursor = self.db.conn.execute("PRAGMA journal_mode")
+        assert cursor.fetchone()[0] == 'wal'
+
+    def test_schema_version(self):
+        cursor = self.db.conn.execute("SELECT MAX(version) FROM schema_version")
+        version = cursor.fetchone()[0]
+        assert version >= 1
+
+    def test_insert_and_get_document(self):
+        from storage.models import Document
+        doc = Document(
+            id="test_hash_123",
+            source="vault7",
+            source_url="https://example.com/test",
+            doc_type="html",
+            title="Test Document",
+            text="This is test content",
+            content_hash="test_hash_123",
+            status="CACHED",
+        )
+        is_new = self.db.insert_document(doc)
+        assert is_new is True
+
+        retrieved = self.db.get_document("test_hash_123")
+        assert retrieved is not None
+        assert retrieved.title == "Test Document"
+        assert retrieved.source == "vault7"
+
+    def test_document_update(self):
+        from storage.models import Document
+        doc = Document(id="up_123", source="test", source_url="http://x",
+                      doc_type="html", title="Original", content_hash="up_123",
+                      status="PENDING")
+        self.db.insert_document(doc)
+
+        doc.title = "Updated"
+        doc.status = "CACHED"
+        is_new = self.db.insert_document(doc)
+        assert is_new is False
+
+        retrieved = self.db.get_document("up_123")
+        assert retrieved.title == "Updated"
+        assert retrieved.status == "CACHED"
+
+    def test_fts_search(self):
+        from storage.models import Document
+        doc = Document(id="fts_123", source="vault7", source_url="http://x",
+                      doc_type="html", title="CIA Hacking Tools Overview",
+                      text="This document describes various CIA hacking tools including HIVE and Grasshopper",
+                      content_hash="fts_123", status="CACHED")
+        self.db.insert_document(doc)
+
+        results = self.db.search_documents("CIA hacking tools")
+        assert len(results) > 0
+        assert results[0].id == "fts_123"
+
+    def test_fts_sanitization(self):
+        # Should not crash on special characters
+        results = self.db.search_documents('test "injection" (attack)')
+        assert isinstance(results, list)
+
+    def test_tool_dedup(self):
+        from storage.models import ExtractedTool
+        tool1 = ExtractedTool(name="HIVE", capability="implant",
+                             description="First", source_doc_ids=["doc1"])
+        tool1.compute_dedup_key()
+        assert self.db.insert_extracted_tool(tool1) is True
+
+        tool2 = ExtractedTool(name="HIVE", capability="implant",
+                             description="Second", source_doc_ids=["doc2"])
+        tool2.compute_dedup_key()
+        assert self.db.insert_extracted_tool(tool2) is False  # Dedup
+
+        # Verify source_doc_ids merged
+        tools = self.db.get_all_tools()
+        assert len(tools) == 1
+        assert "doc1" in tools[0].source_doc_ids
+        assert "doc2" in tools[0].source_doc_ids
+
+    def test_extraction_status(self):
+        self.db.set_extraction_status("doc1", "ttp", "IN_PROGRESS")
+        status = self.db.get_extraction_status("doc1", "ttp")
+        assert status.status == "IN_PROGRESS"
+
+        self.db.set_extraction_status("doc1", "ttp", "DONE", token_count=1500)
+        status = self.db.get_extraction_status("doc1", "ttp")
+        assert status.status == "DONE"
+        assert status.token_count == 1500
+
+    def test_crash_recovery(self):
+        """IN_PROGRESS states should reset to PENDING on DB init."""
+        from storage.models import Document
+        doc = Document(id="crash_test", source="test", source_url="http://x",
+                      doc_type="html", title="Crash", content_hash="crash_test",
+                      status="DOWNLOADING", parse_status="IN_PROGRESS")
+        self.db.insert_document(doc)
+        self.db.set_extraction_status("crash_test", "ttp", "IN_PROGRESS")
+
+        # Re-open database (simulates restart)
+        self.db.close()
+        from storage.db import MosaicDB
+        self.db = MosaicDB(self.tmp)
+
+        doc = self.db.get_document("crash_test")
+        assert doc.status == "PENDING"
+        assert doc.parse_status == "PENDING"
+
+        status = self.db.get_extraction_status("crash_test", "ttp")
+        assert status.status == "PENDING"
+
+    def test_status_counts(self):
+        from storage.models import Document
+        for i, status in enumerate(['CACHED', 'CACHED', 'DOWNLOAD_FAILED']):
+            doc = Document(id=f"sc_{i}", source="test", source_url=f"http://x/{i}",
+                          doc_type="html", title=f"Doc {i}", content_hash=f"sc_{i}",
+                          status=status)
+            self.db.insert_document(doc)
+
+        counts = self.db.get_status_counts("test")
+        assert counts['total'] == 3
+        assert counts['collection']['CACHED'] == 2
+        assert counts['collection']['DOWNLOAD_FAILED'] == 1
+
+    def test_token_usage(self):
+        from storage.models import TokenUsage
+        usage = TokenUsage(
+            timestamp="2024-01-01T00:00:00",
+            source="vault7",
+            doc_id="doc1",
+            extractor="ttp",
+            input_tokens=1000,
+            output_tokens=500,
+            model="claude-sonnet-4-20250514",
+        )
+        self.db.log_token_usage(usage)
+        totals = self.db.get_total_tokens("vault7")
+        assert totals['input_tokens'] == 1000
+        assert totals['output_tokens'] == 500
+        assert totals['total_tokens'] == 1500
+
+
+class TestSanitize:
+    """Test input sanitization utilities."""
+
+    def test_sanitize_filename_basic(self):
+        from utils.sanitize import sanitize_filename
+        assert sanitize_filename("hello.txt") == "hello.txt"
+        assert sanitize_filename("my file (1).pdf") == "my_file_1_.pdf"
+
+    def test_sanitize_filename_traversal(self):
+        from utils.sanitize import sanitize_filename
+        assert sanitize_filename("../../etc/passwd") == "passwd"
+        assert sanitize_filename("/tmp/evil.sh") == "evil.sh"
+
+    def test_sanitize_filename_empty(self):
+        from utils.sanitize import sanitize_filename
+        assert sanitize_filename("") == "unnamed"
+        assert sanitize_filename("...") == "..."
+
+    def test_sanitize_path_safe(self):
+        from utils.sanitize import sanitize_path
+        result = sanitize_path("subdir/file.txt", "/tmp/base")
+        assert result is not None
+        assert result.startswith("/tmp/base")
+
+    def test_sanitize_path_traversal(self):
+        from utils.sanitize import sanitize_path
+        result = sanitize_path("../../etc/passwd", "/tmp/base")
+        assert result is None
+
+    def test_strip_ansi(self):
+        from utils.sanitize import strip_ansi
+        assert strip_ansi("\x1b[31mred\x1b[0m") == "red"
+        assert strip_ansi("no ansi") == "no ansi"
+
+    def test_content_type_fallback(self):
+        """Test magic bytes detection fallback."""
+        from utils.sanitize import _detect_by_magic_bytes
+        # Create a temp PDF-like file
+        with tempfile.NamedTemporaryFile(suffix='.pdf', delete=False) as f:
+            f.write(b'%PDF-1.4 test')
+            f.flush()
+            detected = _detect_by_magic_bytes(f.name)
+            assert detected == 'application/pdf'
+            os.unlink(f.name)
+
+
+class TestTextUtils:
+    """Test text chunking and cleaning."""
+
+    def test_estimate_tokens(self):
+        from utils.text_utils import estimate_tokens
+        assert estimate_tokens("hello world") > 0
+        assert estimate_tokens("") == 0
+
+    def test_clean_text(self):
+        from utils.text_utils import clean_text
+        assert clean_text("hello\x00world") == "hello world"  # null bytes removed (control chars)
+        assert clean_text("line1\r\nline2") == "line1\nline2"
+        assert clean_text("a\n\n\n\nb") == "a\n\nb"
+
+    def test_chunk_default_small(self):
+        from utils.text_utils import chunk_text
+        chunks = chunk_text("short text", doc_type="text")
+        assert len(chunks) == 1
+        assert chunks[0]['text'] == "short text"
+        assert chunks[0]['chunk_index'] == 0
+        assert chunks[0]['total_chunks'] == 1
+
+    def test_chunk_default_large(self):
+        from utils.text_utils import chunk_text
+        # Create text that's definitely larger than default chunk size
+        big_text = "This is a paragraph.\n\n" * 5000
+        chunks = chunk_text(big_text, doc_type="text", max_tokens=1000)
+        assert len(chunks) > 1
+        for i, chunk in enumerate(chunks):
+            assert chunk['chunk_index'] == i
+            assert chunk['total_chunks'] == len(chunks)
+
+    def test_chunk_cable(self):
+        from utils.text_utils import chunk_text
+        text = "SUBJECT: Test Cable\n\nBody paragraph about intelligence."
+        metadata = {'SUBJECT': 'Test Cable', 'CLASSIFICATION': 'CONFIDENTIAL'}
+        chunks = chunk_text(text, doc_type="cable", metadata=metadata)
+        assert len(chunks) >= 1
+        assert 'SUBJECT: Test Cable' in chunks[0]['text']
+
+    def test_dedup_text(self):
+        from utils.text_utils import dedup_text
+        texts = ["hello", "world", "hello", "hello"]
+        unique = dedup_text(texts)
+        assert len(unique) == 2
+
+    def test_chunk_empty(self):
+        from utils.text_utils import chunk_text
+        assert chunk_text("", doc_type="text") == []
+        assert chunk_text("", doc_type="pdf") == []
+
+
+class TestCache:
+    """Test download cache operations."""
+
+    def setup_method(self):
+        self.tmpdir = tempfile.mkdtemp()
+        from utils.cache import DownloadCache
+        self.cache = DownloadCache(self.tmpdir)
+
+    def test_hash_content(self):
+        h = self.cache.hash_content(b"hello world")
+        assert len(h) == 64  # SHA256 hex
+
+    def test_shard_path(self):
+        path = self.cache.get_shard_path("abcdef1234567890", "vault7", "html")
+        assert "vault7" in str(path)
+        assert "/ab/cd/" in str(path)
+        assert path.name == "abcdef1234567890.html"
+
+    def test_store_and_verify(self):
+        content = b"test document content"
+        content_hash = self.cache.hash_content(content)
+
+        # Write temp file
+        tmp_path = os.path.join(self.tmpdir, "test.tmp")
+        with open(tmp_path, 'wb') as f:
+            f.write(content)
+
+        # Store atomically
+        final = self.cache.store_atomic(tmp_path, content_hash, "test_collection", "txt")
+        assert os.path.exists(final)
+        assert not os.path.exists(tmp_path)  # Temp file moved
+
+        # Verify
+        assert self.cache.is_cached(content_hash, "test_collection", "txt")
+        assert self.cache.verify_file(final, content_hash)
+
+    def test_store_hash_mismatch(self):
+        content = b"test content"
+        tmp_path = os.path.join(self.tmpdir, "bad.tmp")
+        with open(tmp_path, 'wb') as f:
+            f.write(content)
+
+        with pytest.raises(ValueError, match="Hash mismatch"):
+            self.cache.store_atomic(tmp_path, "wrong_hash", "test", "txt")
+
+    def test_disk_space_check(self):
+        assert self.cache.check_disk_space(required_mb=1) is True
+
+
+class TestRateLimiters:
+    """Test rate limiting utilities."""
+
+    def test_http_limiter_basic(self):
+        from utils.http_limiter import HTTPRateLimiter
+        limiter = HTTPRateLimiter(rate=100, jitter=0)  # Fast for testing
+        limiter.wait()  # Should not raise
+
+    def test_api_limiter_success(self):
+        from utils.api_limiter import APIRateLimiter
+        limiter = APIRateLimiter()
+        limiter.record_success()
+        assert not limiter.is_circuit_open
+
+    def test_api_limiter_circuit_breaker(self):
+        from utils.api_limiter import APIRateLimiter, CircuitBreakerOpen
+        limiter = APIRateLimiter(circuit_threshold=3)
+        limiter.record_failure("err1")
+        limiter.record_failure("err2")
+        limiter.record_failure("err3")
+        assert limiter.is_circuit_open
+
+        with pytest.raises(CircuitBreakerOpen):
+            limiter.wait_if_needed()
+
+    def test_api_limiter_reset(self):
+        from utils.api_limiter import APIRateLimiter
+        limiter = APIRateLimiter(circuit_threshold=2)
+        limiter.record_failure("err")
+        limiter.record_failure("err")
+        assert limiter.is_circuit_open
+        limiter.reset()
+        assert not limiter.is_circuit_open
+
+
+class TestProfileLoader:
+    """Test YAML profile loading."""
+
+    def test_load_vault7(self):
+        from collectors.profile_loader import load_profile
+        profile = load_profile('vault7')
+        assert profile['name'] == 'vault7'
+        assert 'wikileaks.org' in profile['domain_allowlist']
+        assert profile['crawl_rules']['max_depth'] > 0
+
+    def test_load_nonexistent(self):
+        from collectors.profile_loader import load_profile
+        with pytest.raises(FileNotFoundError):
+            load_profile('nonexistent_profile_xyz')
+
+    def test_list_profiles(self):
+        from collectors.profile_loader import list_profiles
+        profiles = list_profiles()
+        assert len(profiles) >= 5  # vault7, vault8, cablegate, spyfiles, emails
+        names = [p['name'] for p in profiles]
+        assert 'vault7' in names
+
+    def test_create_adhoc(self):
+        from collectors.profile_loader import create_adhoc_profile
+        profile = create_adhoc_profile("https://example.com/docs/", "test_source")
+        assert profile['name'] == 'test_source'
+        assert 'example.com' in profile['domain_allowlist']
+
+
+class TestJSONLStore:
+    """Test JSONL export/import."""
+
+    def setup_method(self):
+        self.tmp_db = tempfile.mktemp(suffix='.db')
+        self.tmp_dir = tempfile.mkdtemp()
+        from storage.db import MosaicDB
+        from storage.jsonl_store import JSONLStore
+        self.db = MosaicDB(self.tmp_db)
+        self.store = JSONLStore(self.db, self.tmp_dir)
+
+    def teardown_method(self):
+        self.db.close()
+        for ext in ('', '-wal', '-shm'):
+            path = self.tmp_db + ext
+            if os.path.exists(path):
+                os.unlink(path)
+
+    def test_export_empty(self):
+        files = self.store.export(category='tools')
+        assert files == []
+
+    def test_export_and_checksum(self):
+        from storage.models import ExtractedTool
+        tool = ExtractedTool(name="TestTool", capability="exploit",
+                           description="A test", source_doc_ids=["doc1"])
+        tool.compute_dedup_key()
+        self.db.insert_extracted_tool(tool)
+
+        files = self.store.export(category='tools')
+        assert len(files) == 1
+        assert files[0].endswith('.jsonl')
+
+        # Verify checksum
+        assert self.store.verify_checksum(files[0])
+
+
+if __name__ == '__main__':
+    pytest.main([__file__, '-v'])