security_demo/test_graph_implementation.py

#!/usr/bin/env python3
"""
Test script for graph implementation
"""
import asyncio
import sys
from pathlib import Path

# Add backend to Python path
backend_path = Path(__file__).parent / "backend"
sys.path.insert(0, str(backend_path))

from app.services.graph_analyzer import graph_analyzer
from app.models.entity import Entity, EntityRelation
from datetime import datetime
import networkx as nx

async def test_graph_implementation():
    """Test the graph implementation with mock data"""
    print("🧪 Testing Graph Implementation")
    print("=" * 40)
    
    # Create mock entities
    entities = []
    base_time = datetime.utcnow()
    
    # Mock entities representing a typical security scenario
    entity_data = [
        {"id": "1", "type": "ip", "value": "192.168.1.100", "name": "Internal Host", "risk": 0.8},
        {"id": "2", "type": "account_id", "value": "admin", "name": "Admin User", "risk": 0.9},
        {"id": "3", "type": "host_id", "value": "workstation-01", "name": "User Workstation", "risk": 0.6},
        {"id": "4", "type": "file_hash", "value": "abc123", "name": "suspicious.exe", "risk": 0.95},
        {"id": "5", "type": "ip", "value": "malicious.com", "name": "External C2", "risk": 1.0},
    ]
    
    for data in entity_data:
        entity = Entity(
            id=data["id"],
            alert_id="test-alert",
            entity_type=data["type"],
            entity_name=data["name"],
            entity_value=data["value"],
            risk_score=data["risk"],
            alert_count=1,
            first_seen=base_time,
            last_seen=base_time
        )
        entities.append(entity)
    
    print(f"✅ Created {len(entities)} mock entities")
    
    # Create mock relations
    relations = []
    relation_data = [
        {"source": "2", "target": "1", "type": "account_from_ip", "strength": 0.8},
        {"source": "2", "target": "3", "type": "account_on_host", "strength": 0.9},
        {"source": "3", "target": "4", "type": "host_contains_file", "strength": 0.7},
        {"source": "1", "target": "5", "type": "ip_communicated_with", "strength": 0.95},
    ]
    
    for data in relation_data:
        relation = EntityRelation(
            source_entity_id=data["source"],
            target_entity_id=data["target"],
            relation_type=data["type"],
            relation_strength=data["strength"],
            evidence_count=1,
            first_observed=base_time,
            last_observed=base_time
        )
        relations.append(relation)
    
    print(f"✅ Created {len(relations)} mock relations")
    
    # Test graph construction
    print("\n🔧 Testing Graph Construction...")
    try:
        # Note: We pass None for db since we're using mock data
        graph = await graph_analyzer.build_entity_graph(None, entities, relations)
        print(f"✅ Graph built successfully: {len(graph.nodes)} nodes, {len(graph.edges)} edges")
        
        # Test centrality analysis
        print("\n📊 Testing Centrality Analysis...")
        centrality_scores = await graph_analyzer.analyze_entity_centrality(graph)
        print("✅ Centrality scores calculated:")
        for node_id, score in sorted(centrality_scores.items(), key=lambda x: x[1], reverse=True):
            entity = next((e for e in entities if str(e.id) == node_id), None)
            entity_name = entity.entity_name if entity else "Unknown"
            print(f"   - {entity_name} ({node_id}): {score:.3f}")
        
        # Test attack chain discovery
        print("\n🔗 Testing Attack Chain Discovery...")
        attack_chains = await graph_analyzer.find_attack_chains(graph, max_chain_length=5)
        print(f"✅ Found {len(attack_chains)} attack chains:")
        for i, chain in enumerate(attack_chains[:3]):
            chain_names = []
            for node_id in chain:
                entity = next((e for e in entities if str(e.id) == node_id), None)
                entity_name = entity.entity_name if entity else "Unknown"
                chain_names.append(f"{entity_name}({node_id})")
            print(f"   Chain {i+1}: {' → '.join(chain_names)}")
        
        # Test subgraph discovery
        print("\n🗂️  Testing Subgraph Discovery...")
        subgraphs = await graph_analyzer.find_subgraphs(graph, min_size=2)
        print(f"✅ Found {len(subgraphs)} subgraphs:")
        for i, subgraph in enumerate(subgraphs):
            importance = subgraph.graph.get('importance', 0)
            print(f"   Subgraph {i+1}: {len(subgraph.nodes)} nodes, importance: {importance:.3f}")
        
        print("\n🎉 All graph functionality tests passed!")
        return True
        
    except Exception as e:
        print(f"❌ Graph test failed: {e}")
        import traceback
        traceback.print_exc()
        return False

if __name__ == "__main__":
    success = asyncio.run(test_graph_implementation())
    if success:
        print("\n✅ Graph implementation is working correctly!")
    else:
        print("\n❌ Graph implementation has issues!")
        sys.exit(1)