e2e_test.py

#!/usr/bin/env python3
"""
EdgeAI Telemetry - E2E Test Suite
Simulates the full data flow without Docker
"""

import json
import time
import hashlib
import random
from datetime import datetime, timedelta, timezone
from dataclasses import dataclass, asdict
from typing import List, Dict, Optional

# Test Configuration
TEST_CONFIG = {
    "total_events": 10000,
    "window_seconds": 60,
    "identities": [
        "user1@example.com",
        "user2@example.com", 
        "admin@example.com",
        "attacker@external.com",
        "service-account"
    ],
    "event_types": [
        "login_success",
        "login_failure",
        "file_access",
        "db_query",
        "api_request",
        "privilege_escalation"
    ]
}

@dataclass
class RawEvent:
    """Event as generated by source"""
    timestamp: str
    identity: str
    event_type: str
    severity: str
    metadata: dict

@dataclass
class EventSummary:
    """Aggregated summary (what edge agent sends)"""
    identity_hash: str
    event_type_id: int
    count: int
    window_start: int
    window_end: int
    severity_dist: dict
    baseline_score: float
    agent_id: str

@dataclass
class AnomalyEvent:
    """Anomaly flagged by edge detection"""
    identity_hash: str
    event_type_id: int
    timestamp: int
    severity: int
    detection_reason: str
    anomaly_score: float

class EdgeAgentSimulator:
    """Simulates edge agent aggregation logic"""
    
    def __init__(self, window_seconds: int = 60):
        self.window_seconds = window_seconds
        self.event_type_map: Dict[str, int] = {}
        self.next_type_id = 1
        self.windows: Dict[tuple, dict] = {}
        self.agent_id = "test-agent-001"
        
    def _hash_identity(self, identity: str) -> str:
        """Hash identity to 16-byte string"""
        return hashlib.blake2b(identity.encode(), digest_size=16).hexdigest()[:32]
    
    def _get_type_id(self, event_type: str) -> int:
        """Map event type to integer ID"""
        if event_type not in self.event_type_map:
            self.event_type_map[event_type] = self.next_type_id
            self.next_type_id += 1
        return self.event_type_map[event_type]
    
    def _severity_to_num(self, severity: str) -> int:
        """Convert severity string to number"""
        mapping = {"critical": 5, "high": 4, "medium": 3, "low": 2, "info": 1}
        return mapping.get(severity.lower(), 1)
    
    def _calculate_baseline(self, count: int, severities: dict) -> float:
        """Calculate anomaly baseline score"""
        base_score = min(count / 100.0, 1.0) * 20.0
        severity_weight = sum(sev * cnt * 0.1 for sev, cnt in severities.items())
        return min(base_score + severity_weight, 100.0)
    
    def _check_anomaly(self, identity: str, event: RawEvent, count: int) -> Optional[AnomalyEvent]:
        """Simple anomaly detection rules"""
        score = 0.0
        reasons = []
        
        # Rule 1: High volume (brute force)
        if count >= 5:
            score += 40.0
            reasons.append("volume_spike")
        
        # Rule 2: High severity
        if event.severity in ["critical", "high"]:
            score += 30.0
            reasons.append("high_severity")
        
        # Rule 3: Suspicious identity
        if "attacker" in identity or "admin" in identity:
            score += 20.0
            reasons.append("sensitive_identity")
        
        if score >= 70.0:
            return AnomalyEvent(
                identity_hash=self._hash_identity(identity),
                event_type_id=self._get_type_id(event.event_type),
                timestamp=int(datetime.fromisoformat(event.timestamp.replace('Z', '+00:00')).timestamp()),
                severity=self._severity_to_num(event.severity),
                detection_reason=",".join(reasons),
                anomaly_score=score
            )
        return None
    
    def process_events(self, events: List[RawEvent]) -> tuple:
        """Process events through aggregation windows"""
        summaries = []
        anomalies = []
        
        # Group by window and identity+type
        windows: Dict[tuple, dict] = {}
        
        for event in events:
            ts = datetime.fromisoformat(event.timestamp.replace('Z', '+00:00'))
            window_start = int(ts.timestamp()) // self.window_seconds * self.window_seconds
            
            key = (event.identity, event.event_type, window_start)
            
            if key not in windows:
                windows[key] = {
                    "identity": event.identity,
                    "event_type": event.event_type,
                    "count": 0,
                    "severities": {},
                    "window_start": window_start,
                    "window_end": window_start + self.window_seconds
                }
            
            window = windows[key]
            window["count"] += 1
            sev_num = self._severity_to_num(event.severity)
            window["severities"][sev_num] = window["severities"].get(sev_num, 0) + 1
            
            # Check for anomaly
            if anomaly := self._check_anomaly(event.identity, event, window["count"]):
                if anomaly not in anomalies:
                    anomalies.append(anomaly)
        
        # Convert windows to summaries
        for window in windows.values():
            summary = EventSummary(
                identity_hash=self._hash_identity(window["identity"]),
                event_type_id=self._get_type_id(window["event_type"]),
                count=window["count"],
                window_start=window["window_start"],
                window_end=window["window_end"],
                severity_dist=window["severities"],
                baseline_score=self._calculate_baseline(window["count"], window["severities"]),
                agent_id=self.agent_id
            )
            summaries.append(summary)
        
        return summaries, anomalies

class CloudServerSimulator:
    """Simulates cloud server correlation"""
    
    def __init__(self):
        self.summaries: List[EventSummary] = []
        self.anomalies: List[AnomalyEvent] = []
        self.incidents: List[dict] = []
        
    def receive_data(self, summaries: List[EventSummary], anomalies: List[AnomalyEvent]):
        """Receive aggregated data from edge"""
        self.summaries.extend(summaries)
        self.anomalies.extend(anomalies)
        
        # Correlate anomalies into incidents
        self._correlate_anomalies()
    
    def _correlate_anomalies(self):
        """Simple correlation logic"""
        # Group anomalies by identity hash
        by_identity: Dict[str, List[AnomalyEvent]] = {}
        for anomaly in self.anomalies:
            by_identity.setdefault(anomaly.identity_hash, []).append(anomaly)
        
        for identity_hash, anoms in by_identity.items():
            # Pattern: Brute force (5+ auth failures)
            auth_failures = [a for a in anoms if a.anomaly_score > 40]
            if len(auth_failures) >= 3:
                self.incidents.append({
                    "id": f"inc-{len(self.incidents)}",
                    "identity_hash": identity_hash,
                    "pattern": "brute_force_attempt",
                    "risk_score": 85.0,
                    "anomaly_count": len(auth_failures)
                })
            
            # Pattern: Severity burst (3+ high severity)
            high_sev = [a for a in anoms if a.severity >= 4]
            if len(high_sev) >= 2:
                self.incidents.append({
                    "id": f"inc-{len(self.incidents)}",
                    "identity_hash": identity_hash,
                    "pattern": "severity_burst",
                    "risk_score": 75.0,
                    "anomaly_count": len(high_sev)
                })

def generate_test_events(count: int) -> List[RawEvent]:
    """Generate realistic test events"""
    events = []
    base_time = datetime.now(timezone.utc)
    
    for i in range(count):
        # Create patterns
        if i < 100:  # Normal user activity
            identity = random.choice(["user1@example.com", "user2@example.com"])
            event_type = random.choice(["login_success", "file_access", "api_request"])
            severity = "info"
        elif i < 150:  # Brute force attack
            identity = "attacker@external.com"
            event_type = "login_failure"
            severity = "high"
        elif i < 160:  # Privilege escalation
            identity = "admin@example.com"
            event_type = "privilege_escalation"
            severity = "critical"
        else:  # Mixed
            identity = random.choice(TEST_CONFIG["identities"])
            event_type = random.choice(TEST_CONFIG["event_types"])
            severity = random.choice(["info", "low", "medium"])
        
        event = RawEvent(
            timestamp=(base_time + timedelta(seconds=i)).isoformat().replace('+00:00', 'Z'),
            identity=identity,
            event_type=event_type,
            severity=severity,
            metadata={"ip": f"192.168.1.{random.randint(1, 255)}", "index": i}
        )
        events.append(event)
    
    return events

def calculate_data_reduction(raw_events: int, summaries: int) -> float:
    """Calculate percentage reduction"""
    if raw_events == 0:
        return 0.0
    return ((raw_events - summaries) / raw_events) * 100

def main():
    print("=" * 70)
    print("EDGEAI TELEMETRY - E2E TEST SUITE")
    print("=" * 70)
    print()
    
    # Phase 1: Generate Events
    print("[1] PHASE 1: Event Generation")
    print("-" * 70)
    print(f"Generating {TEST_CONFIG['total_events']:,} raw events...")
    
    events = generate_test_events(TEST_CONFIG["total_events"])
    
    # Calculate raw data size
    raw_json = json.dumps([asdict(e) for e in events])
    raw_size_bytes = len(raw_json.encode('utf-8'))
    raw_size_mb = raw_size_bytes / (1024 * 1024)
    
    print(f"[OK] Generated {len(events):,} events")
    print(f"[OK] Raw data size: {raw_size_mb:.2f} MB")
    print(f"[OK] Sample event: {json.dumps(asdict(events[0]), indent=2)}")
    print()
    
    # Phase 2: Edge Aggregation
    print("[2] PHASE 2: Edge Agent Aggregation")
    print("-" * 70)
    
    agent = EdgeAgentSimulator(window_seconds=60)
    start_time = time.time()
    
    summaries, anomalies = agent.process_events(events)
    
    edge_time = time.time() - start_time
    
    # Calculate summary size (simulated protobuf + zstd)
    summary_json = json.dumps([asdict(s) for s in summaries])
    summary_size_bytes = len(summary_json.encode('utf-8'))
    # Apply compression factor (protobuf ~50% + zstd 5:1)
    compressed_size_bytes = summary_size_bytes * 0.5 * 0.2
    compressed_size_mb = compressed_size_bytes / (1024 * 1024)
    
    print(f"[OK] Aggregation completed in {edge_time:.3f}s")
    print(f"[OK] Summaries created: {len(summaries)}")
    print(f"[OK] Anomalies detected: {len(anomalies)}")
    print(f"[OK] Unique identities: {len(set(s.identity_hash for s in summaries))}")
    print(f"[OK] Compressed size: {compressed_size_mb:.4f} MB")
    print()
    
    # Show sample summary
    if summaries:
        print("Sample Summary:")
        print(json.dumps(asdict(summaries[0]), indent=2))
        print()
    
    # Show anomalies
    if anomalies:
        print("Detected Anomalies:")
        for a in anomalies[:3]:
            print(f"  - Score: {a.anomaly_score:.1f}, Reason: {a.detection_reason}")
        print()
    
    # Phase 3: Cloud Correlation
    print("[3] PHASE 3: Cloud Server Correlation")
    print("-" * 70)
    
    cloud = CloudServerSimulator()
    start_time = time.time()
    
    cloud.receive_data(summaries, anomalies)
    
    cloud_time = time.time() - start_time
    
    print(f"[OK] Correlation completed in {cloud_time:.3f}s")
    print(f"[OK] Total summaries stored: {len(cloud.summaries)}")
    print(f"[OK] Total anomalies stored: {len(cloud.anomalies)}")
    print(f"[OK] Incidents created: {len(cloud.incidents)}")
    print()
    
    if cloud.incidents:
        print("Created Incidents:")
        for inc in cloud.incidents:
            print(f"  - {inc['pattern']}: Risk={inc['risk_score']}, Count={inc['anomaly_count']}")
        print()
    
    # Phase 4: Metrics & Validation
    print("[4] PHASE 4: Results & Cost Analysis")
    print("-" * 70)
    
    reduction_pct = calculate_data_reduction(len(events), len(summaries))
    compression_ratio = raw_size_mb / compressed_size_mb if compressed_size_mb > 0 else 0
    
    print(f"{'Metric':<30} {'Value':<20}")
    print("-" * 50)
    print(f"{'Raw Events':<30} {len(events):>20,}")
    print(f"{'Summaries Sent':<30} {len(summaries):>20,}")
    print(f"{'Data Reduction':<30} {reduction_pct:>19.1f}%")
    print(f"{'Compression Ratio':<30} {compression_ratio:>19.1f}:1")
    print(f"{'Anomalies Detected':<30} {len(anomalies):>20}")
    print(f"{'Incidents Correlated':<30} {len(cloud.incidents):>20}")
    print()
    
    # Cost Analysis
    print("[$] COST ANALYSIS (Monthly, 10k EPS)")
    print("-" * 70)
    
    # Traditional SIEM costs
    traditional_network = 5000  # $/mo
    traditional_storage = 10000  # $/mo
    traditional_compute = 3000   # $/mo
    traditional_total = traditional_network + traditional_storage + traditional_compute
    
    # EdgeAI costs (proportional to data volume)
    edgeai_network = traditional_network * (1 - reduction_pct/100)
    edgeai_storage = traditional_storage * (len(summaries) / len(events))
    edgeai_compute = traditional_compute * 0.1  # Only correlation
    edgeai_total = edgeai_network + edgeai_storage + edgeai_compute
    
    savings = traditional_total - edgeai_total
    
    print(f"{'Component':<25} {'Traditional':>15} {'EdgeAI':>15} {'Savings':>15}")
    print("-" * 70)
    print(f"{'Network Ingress':<25} ${traditional_network:>14,.0f} ${edgeai_network:>14,.0f} ${traditional_network-edgeai_network:>14,.0f}")
    print(f"{'Storage':<25} ${traditional_storage:>14,.0f} ${edgeai_storage:>14,.0f} ${traditional_storage-edgeai_storage:>14,.0f}")
    print(f"{'Compute':<25} ${traditional_compute:>14,.0f} ${edgeai_compute:>14,.0f} ${traditional_compute-edgeai_compute:>14,.0f}")
    print("-" * 70)
    print(f"{'TOTAL':<25} ${traditional_total:>14,.0f} ${edgeai_total:>14,.0f} ${savings:>14,.0f}")
    print()
    
    # Success Criteria
    print("[OK] SUCCESS CRITERIA")
    print("-" * 70)
    
    checks = [
        ("Data Reduction >= 95%", reduction_pct >= 95),
        ("Summaries < 1% of events", len(summaries) < len(events) * 0.01),
        ("Anomalies Detected", len(anomalies) > 0),
        ("Incidents Correlated", len(cloud.incidents) > 0),
        ("Processing Time < 5s", edge_time + cloud_time < 5),
    ]
    
    all_passed = True
    for check_name, passed in checks:
        status = "[OK] PASS" if passed else "[XX] FAIL"
        print(f"  {status}: {check_name}")
        if not passed:
            all_passed = False
    
    print()
    print("=" * 70)
    if all_passed:
        print("[SUCCESS] ALL TESTS PASSED - System is production ready!")
    else:
        print("[WARNING] SOME TESTS FAILED - Review output above")
    print("=" * 70)
    
    return all_passed

if __name__ == "__main__":
    success = main()
    exit(0 if success else 1)