security-platform/standalone_graph_test.py

#!/usr/bin/env python3
"""
Standalone graph functionality test
Tests the core NetworkX-based graph analysis without database dependencies
"""
import asyncio
import sys
from pathlib import Path
from datetime import datetime

# Import networkx directly
import networkx as nx

# Mock the graph analyzer functionality 
class StandaloneGraphAnalyzer:
    """Standalone version of GraphAnalyzer for testing"""
    
    def __init__(self):
        self.max_entities = 1000
        
        # 图分析算法权重
        self.centrality_weights = {
            'degree': 0.3,
            'betweenness': 0.3, 
            'pagerank': 0.4
        }
        
        # 实体类型重要性权重
        self.entity_importance = {
            'account_id': 0.9,
            'ip': 0.8,
            'host_id': 0.7,
            'file_hash': 0.6,
            'asset_group': 0.5
        }
    
    async def build_entity_graph(self, entities, relations):
        """构建实体关系图"""
        graph = nx.Graph()
        
        # 添加节点
        for entity in entities:
            graph.add_node(
                str(entity.id),
                entity_type=entity.entity_type,
                entity_value=entity.entity_value,
                entity_name=entity.entity_name,
                risk_score=entity.risk_score or 0.0,
                alert_count=entity.alert_count,
                first_seen=entity.first_seen,
                last_seen=entity.last_seen
            )
        
        # 添加边
        for relation in relations:
            if (str(relation.source_entity_id) in graph.nodes and 
                str(relation.target_entity_id) in graph.nodes):
                graph.add_edge(
                    str(relation.source_entity_id),
                    str(relation.target_entity_id),
                    relation_type=relation.relation_type,
                    strength=relation.relation_strength,
                    evidence_count=relation.evidence_count,
                    first_observed=relation.first_observed,
                    last_observed=relation.last_observed
                )
        
        return graph
    
    async def analyze_entity_centrality(self, graph):
        """分析实体的中心性指标"""
        if len(graph.nodes) == 0:
            return {}
        
        centrality_scores = {}
        
        try:
            # 度中心性
            degree_centrality = nx.degree_centrality(graph)
            
            # 介数中心性
            betweenness_centrality = nx.betweenness_centrality(graph)
            
            # PageRank
            pagerank = nx.pagerank(graph, weight='strength')
            
            # 综合计算中心性分数
            for node in graph.nodes():
                score = (
                    degree_centrality.get(node, 0) * self.centrality_weights['degree'] +
                    betweenness_centrality.get(node, 0) * self.centrality_weights['betweenness'] +
                    pagerank.get(node, 0) * self.centrality_weights['pagerank']
                )
                centrality_scores[node] = score
        except Exception as e:
            print(f"中心性分析失败: {e}")
            # 降级到简单度计算
            for node in graph.nodes():
                centrality_scores[node] = graph.degree(node) / len(graph.nodes)
        
        return centrality_scores
    
    async def find_attack_chains(self, graph, max_chain_length=6):
        """发现攻击链路径"""
        attack_chains = []
        
        # 查找高风险节点作为起点
        high_risk_nodes = [
            node for node, attrs in graph.nodes(data=True)
            if attrs.get('risk_score', 0) > 0.7
        ]
        
        # 从每个高风险节点开始搜索路径
        for start_node in high_risk_nodes:
            chains = self._dfs_attack_chains(
                graph, start_node, max_chain_length
            )
            attack_chains.extend(chains)
        
        # 按路径风险分数排序
        scored_chains = []
        for chain in attack_chains:
            score = await self._score_attack_chain(graph, chain)
            scored_chains.append((chain, score))
        
        scored_chains.sort(key=lambda x: x[1], reverse=True)
        return [chain for chain, score in scored_chains[:10]]
    
    def _dfs_attack_chains(self, graph, start_node, max_length, current_path=None, visited=None):
        """深度优先搜索攻击链"""
        if current_path is None:
            current_path = [start_node]
        if visited is None:
            visited = {start_node}
        
        chains = []
        
        if len(current_path) >= 3:
            chains.append(current_path.copy())
        
        if len(current_path) >= max_length:
            return chains
        
        # 继续搜索邻居节点
        for neighbor in graph.neighbors(start_node):
            if neighbor not in visited:
                if self._is_valid_chain_step(graph, start_node, neighbor):
                    visited.add(neighbor)
                    current_path.append(neighbor)
                    
                    chains.extend(
                        self._dfs_attack_chains(
                            graph, neighbor, max_length, current_path, visited
                        )
                    )
                    
                    current_path.pop()
                    visited.remove(neighbor)
        
        return chains
    
    def _is_valid_chain_step(self, graph, current_node, next_node):
        """检查是否是有效的攻击链步骤"""
        edge_attrs = graph.edges[current_node, next_node]
        
        # 关系强度检查
        if edge_attrs.get('strength', 0) < 0.3:
            return False
        
        return True
    
    async def _score_attack_chain(self, graph, chain):
        """计算攻击链风险分数"""
        if len(chain) < 2:
            return 0.0
        
        total_score = 0.0
        
        # 节点风险分数
        node_scores = [
            graph.nodes[node].get('risk_score', 0) for node in chain
        ]
        total_score += sum(node_scores) * 0.4
        
        # 边权重分数
        edge_scores = []
        for i in range(len(chain) - 1):
            edge_data = graph.edges[chain[i], chain[i + 1]]
            edge_scores.append(edge_data.get('strength', 0))
        total_score += sum(edge_scores) * 0.3
        
        # 链路长度奖励
        length_bonus = min(len(chain) / 6.0, 1.0) * 0.2
        total_score += length_bonus
        
        # 实体类型多样性
        entity_types = set(
            graph.nodes[node].get('entity_type') for node in chain
        )
        diversity_bonus = len(entity_types) / len(chain) * 0.1
        total_score += diversity_bonus
        
        return min(total_score, 1.0)
    
    async def find_subgraphs(self, graph, min_size=3):
        """发现重要子图"""
        subgraphs = []
        
        # 找到连通分量
        connected_components = list(nx.connected_components(graph))
        
        for component in connected_components:
            if len(component) >= min_size:
                subgraph = graph.subgraph(component).copy()
                
                # 计算子图重要性
                importance = await self._calculate_subgraph_importance(subgraph)
                subgraph.graph['importance'] = importance
                
                subgraphs.append(subgraph)
        
        # 按重要性排序
        subgraphs.sort(
            key=lambda sg: sg.graph.get('importance', 0), reverse=True
        )
        
        return subgraphs
    
    async def _calculate_subgraph_importance(self, subgraph):
        """计算子图重要性分数"""
        if len(subgraph.nodes) == 0:
            return 0.0
        
        # 节点风险分数
        risk_scores = [
            node_attrs.get('risk_score', 0)
            for _, node_attrs in subgraph.nodes(data=True)
        ]
        avg_risk = sum(risk_scores) / len(risk_scores)
        
        # 节点数量权重
        size_weight = min(len(subgraph.nodes) / 10.0, 1.0)
        
        # 边密度
        density = nx.density(subgraph)
        
        # 实体类型多样性
        entity_types = set(
            node_attrs.get('entity_type')
            for _, node_attrs in subgraph.nodes(data=True)
        )
        type_diversity = len(entity_types) / len(subgraph.nodes)
        
        # 综合计算
        importance = (
            avg_risk * 0.4 +
            size_weight * 0.2 +
            density * 0.2 +
            type_diversity * 0.2
        )
        
        return min(importance, 1.0)

# Mock classes for entities and relations
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

async def test_standalone_graph_functionality():
    """Test standalone graph functionality"""
    print("🧪 Testing Standalone Graph Functionality")
    print("=" * 50)
    
    # Initialize analyzer
    analyzer = StandaloneGraphAnalyzer()
    
    # Create mock entities representing a realistic attack scenario
    base_time = datetime.utcnow()
    entities = [
        MockEntity("1", "ip", "192.168.1.100", "Compromised Workstation", 0.8, 5, base_time, base_time),
        MockEntity("2", "account_id", "jdoe", "John Doe", 0.6, 3, base_time, base_time),
        MockEntity("3", "host_id", "WS-001", "Workstation-001", 0.7, 4, base_time, base_time),
        MockEntity("4", "file_hash", "a1b2c3d4e5", "malware.exe", 0.95, 2, base_time, base_time),
        MockEntity("5", "ip", "45.76.123.45", "C2 Server", 1.0, 1, base_time, base_time),
        MockEntity("6", "account_id", "admin", "Admin Account", 0.9, 2, base_time, base_time),
        MockEntity("7", "host_id", "SRV-001", "Domain Controller", 0.85, 3, base_time, base_time),
    ]
    
    # Create mock relations representing attack progression
    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),
        MockRelation("6", "7", "account_on_host", 0.9, 2, base_time, base_time),
        MockRelation("2", "6", "account_escalated_to", 0.8, 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(entities, relations)
        print(f"✅ Graph built successfully: {len(graph.nodes)} nodes, {len(graph.edges)} edges")
        
        # Verify graph structure
        print("   Graph nodes:")
        for node_id, attrs in graph.nodes(data=True):
            print(f"     - {attrs['entity_name']} ({attrs['entity_type']}): risk={attrs['risk_score']:.2f}")
        
        print("   Graph edges:")
        for source, target, attrs in graph.edges(data=True):
            source_name = graph.nodes[source]['entity_name']
            target_name = graph.nodes[target]['entity_name']
            print(f"     - {source_name} → {target_name} ({attrs['relation_type']}): strength={attrs['strength']:.2f}")
        
    except Exception as e:
        print(f"❌ Graph construction failed: {e}")
        import traceback
        traceback.print_exc()
        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 by centrality
        sorted_centrality = sorted(centrality_scores.items(), key=lambda x: x[1], reverse=True)
        print("   Node centrality scores:")
        for node_id, score in sorted_centrality:
            entity_name = graph.nodes[node_id]['entity_name']
            risk_score = graph.nodes[node_id]['risk_score']
            print(f"     - {entity_name}: centrality={score:.4f}, risk={risk_score:.2f}")
        
    except Exception as e:
        print(f"❌ Centrality analysis failed: {e}")
        import traceback
        traceback.print_exc()
        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=6)
        print(f"✅ Found {len(attack_chains)} attack chains")
        
        for i, chain in enumerate(attack_chains[:5]):  # Show first 5 chains
            chain_info = []
            total_risk = 0.0
            for node_id in chain:
                entity_name = graph.nodes[node_id]['entity_name']
                risk = graph.nodes[node_id]['risk_score']
                total_risk += risk
                chain_info.append(f"{entity_name}({risk:.1f})")
            avg_risk = total_risk / len(chain)
            print(f"   Chain {i+1} (avg risk: {avg_risk:.2f}): {' → '.join(chain_info)}")
        
    except Exception as e:
        print(f"❌ Attack chain discovery failed: {e}")
        import traceback
        traceback.print_exc()
        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_name = graph.nodes[node_id]['entity_name']
                node_names.append(entity_name)
            print(f"   Subgraph {i+1} (importance: {importance:.3f}): [{', '.join(node_names)}]")
        
    except Exception as e:
        print(f"❌ Subgraph discovery failed: {e}")
        import traceback
        traceback.print_exc()
        return False
    
    print("\n🎉 All graph functionality tests passed!")
    print("\n📈 Analysis Summary:")
    print(f"   - Total Entities: {len(entities)}")
    print(f"   - Total Relations: {len(relations)}")
    print(f"   - Graph Nodes: {len(graph.nodes)}")
    print(f"   - Graph Edges: {len(graph.edges)}")
    print(f"   - Attack Chains Found: {len(attack_chains)}")
    print(f"   - Subgraphs Found: {len(subgraphs)}")
    print(f"   - Average Entity Risk: {sum(e.risk_score for e in entities) / len(entities):.2f}")
    print(f"   - High Risk Entities (>0.8): {len([e for e in entities if e.risk_score > 0.8])}")
    
    return True

if __name__ == "__main__":
    try:
        success = asyncio.run(test_standalone_graph_functionality())
        if success:
            print("\n✅ Graph implementation is fully functional!")
            print("🚀 Ready for integration with the backend system!")
        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)