security-platform/simple_graph_test.py

#!/usr/bin/env python3
"""
Simple graph functionality test without database dependencies
"""
import asyncio
import sys
import os
from pathlib import Path

# Add backend to Python path and set up environment
backend_path = Path(__file__).parent / "backend"
sys.path.insert(0, str(backend_path))
os.environ["DATABASE_URL"] = "sqlite+aiosqlite:///test.db"  
os.environ["OPENAI_API_KEY"] = "test-key"

import networkx as nx
from datetime import datetime

# Simple mock classes to avoid database dependencies
class MockEntity:
    def __init__(self, id, entity_type, entity_value, entity_name, risk_score, alert_count, first_seen, last_seen):
        self.id = id
        self.entity_type = entity_type
        self.entity_value = entity_value
        self.entity_name = entity_name
        self.risk_score = risk_score
        self.alert_count = alert_count
        self.first_seen = first_seen
        self.last_seen = last_seen

class MockRelation:
    def __init__(self, source_entity_id, target_entity_id, relation_type, relation_strength, evidence_count, first_observed, last_observed):
        self.source_entity_id = source_entity_id
        self.target_entity_id = target_entity_id
        self.relation_type = relation_type
        self.relation_strength = relation_strength
        self.evidence_count = evidence_count
        self.first_observed = first_observed
        self.last_observed = last_observed

# Import the core graph analyzer functionality
from app.services.graph_analyzer import GraphAnalyzer

async def test_core_graph_functionality():
    """Test core graph functionality"""
    print("🧪 Testing Core Graph Functionality")
    print("=" * 45)
    
    # Initialize analyzer
    analyzer = GraphAnalyzer()
    
    # Create mock entities
    base_time = datetime.utcnow()
    entities = [
        MockEntity("1", "ip", "192.168.1.100", "Internal Host", 0.8, 5, base_time, base_time),
        MockEntity("2", "account_id", "admin", "Admin User", 0.9, 3, base_time, base_time),
        MockEntity("3", "host_id", "workstation-01", "User Workstation", 0.6, 4, base_time, base_time),
        MockEntity("4", "file_hash", "abc123def", "suspicious.exe", 0.95, 2, base_time, base_time),
        MockEntity("5", "ip", "8.8.8.8", "External Server", 1.0, 1, base_time, base_time),
    ]
    
    # Create mock relations
    relations = [
        MockRelation("2", "1", "account_from_ip", 0.8, 2, base_time, base_time),
        MockRelation("2", "3", "account_on_host", 0.9, 3, base_time, base_time),
        MockRelation("3", "4", "host_contains_file", 0.7, 1, base_time, base_time),
        MockRelation("1", "5", "ip_communicated_with", 0.95, 1, base_time, base_time),
    ]
    
    print(f"✅ Created {len(entities)} mock entities")
    print(f"✅ Created {len(relations)} mock relations")
    
    # Test 1: Graph Construction
    print("\n🔧 Test 1: Graph Construction")
    try:
        graph = await analyzer.build_entity_graph(None, entities, relations)  # Pass None for db
        print(f"✅ Graph built: {len(graph.nodes)} nodes, {len(graph.edges)} edges")
        
        # Verify graph structure
        expected_nodes = {str(e.id) for e in entities}
        actual_nodes = set(graph.nodes())
        assert expected_nodes == actual_nodes, f"Node mismatch: expected {expected_nodes}, got {actual_nodes}"
        print("✅ All expected nodes present")
        
        # Verify edges
        print("✅ Graph edges:")
        for source, target, data in graph.edges(data=True):
            print(f"   - {source} → {target} ({data.get('relation_type', 'unknown')})")
        
    except Exception as e:
        print(f"❌ Graph construction failed: {e}")
        return False
    
    # Test 2: Centrality Analysis  
    print("\n📊 Test 2: Centrality Analysis")
    try:
        centrality_scores = await analyzer.analyze_entity_centrality(graph)
        print(f"✅ Centrality analysis completed for {len(centrality_scores)} nodes")
        
        # Display results
        sorted_centrality = sorted(centrality_scores.items(), key=lambda x: x[1], reverse=True)
        for node_id, score in sorted_centrality:
            entity = next((e for e in entities if str(e.id) == node_id), None)
            name = entity.entity_name if entity else "Unknown"
            print(f"   - {name}: {score:.4f}")
        
    except Exception as e:
        print(f"❌ Centrality analysis failed: {e}")
        return False
    
    # Test 3: Attack Chain Discovery
    print("\n🔗 Test 3: Attack Chain Discovery")
    try:
        attack_chains = await 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]):  # Show first 3 chains
            chain_info = []
            for node_id in chain:
                entity = next((e for e in entities if str(e.id) == node_id), None)
                name = entity.entity_name if entity else f"Node({node_id})"
                risk = entity.risk_score if entity else 0.0
                chain_info.append(f"{name}({risk:.1f})")
            print(f"   Chain {i+1}: {' → '.join(chain_info)}")
        
    except Exception as e:
        print(f"❌ Attack chain discovery failed: {e}")
        return False
    
    # Test 4: Subgraph Discovery
    print("\n🗂️  Test 4: Subgraph Discovery")
    try:
        subgraphs = await 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.0)
            node_names = []
            for node_id in subgraph.nodes():
                entity = next((e for e in entities if str(e.id) == node_id), None)
                name = entity.entity_name if entity else f"Node({node_id})"
                node_names.append(name)
            print(f"   Subgraph {i+1}: [{', '.join(node_names)}] (importance: {importance:.3f})")
        
    except Exception as e:
        print(f"❌ Subgraph discovery failed: {e}")
        return False
    
    print("\n🎉 All core graph functionality tests passed!")
    print("\n📈 Summary:")
    print(f"   - Nodes: {len(graph.nodes)}")
    print(f"   - Edges: {len(graph.edges)}")
    print(f"   - Attack Chains: {len(attack_chains)}")
    print(f"   - Subgraphs: {len(subgraphs)}")
    print(f"   - Avg Risk: {sum(e.risk_score for e in entities) / len(entities):.2f}")
    
    return True

if __name__ == "__main__":
    try:
        success = asyncio.run(test_core_graph_functionality())
        if success:
            print("\n✅ Graph implementation is fully functional!")
        else:
            print("\n❌ Graph implementation has issues!")
            sys.exit(1)
    except KeyboardInterrupt:
        print("\n⚠️  Test interrupted by user")
    except Exception as e:
        print(f"\n❌ Unexpected error: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)