"""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'])