n0mad1k
501 lines
18 KiB
Python
501 lines
18 KiB
Python
"""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") == "helloworld" # null bytes stripped
|
|
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'])
|