correlation.py

"""
Correlation engine for detecting coordinated activities and lateral movement.
"""
import logging
from typing import List, Optional, Dict, Set, Tuple
from dataclasses import dataclass
from datetime import datetime, timedelta
from collections import defaultdict

from models import AuditEvent, RiskLevel


@dataclass
class Correlation:
    """Represents correlated events."""
    events: List[AuditEvent]
    correlation_type: str
    description: str
    risk_level: RiskLevel
    timestamp: datetime
    users_involved: Set[Optional[str]]
    auids_involved: Set[Optional[int]]


@dataclass
class LateralMovement:
    """Represents lateral movement pattern."""
    source_user: Optional[str]
    target_user: Optional[str]
    source_auid: Optional[int]
    target_auid: Optional[int]
    events: List[AuditEvent]
    timestamp: datetime
    risk_level: RiskLevel


class CorrelationEngine:
    """Correlate events across users and time."""
    
    # Time window for correlation (seconds)
    CORRELATION_TIME_WINDOW = 300  # 5 minutes
    
    def __init__(self, logger: Optional[logging.Logger] = None):
        """Initialize correlation engine.
        
        Args:
            logger: Optional logger instance
        """
        self.logger = logger or logging.getLogger(__name__)
    
    def detect_coordinated_activity(self, events: List[AuditEvent]) -> List[Correlation]:
        """Detect coordinated suspicious activity.
        
        Args:
            events: List of audit events
            
        Returns:
            List of correlations
        """
        correlations: List[Correlation] = []
        
        # Group events by time windows
        time_windows = self._group_by_time_windows(events)
        
        for window_events in time_windows:
            # Check for similar commands from different users
            similar_commands = self._find_similar_commands(window_events)
            if similar_commands:
                correlations.extend(similar_commands)
            
            # Check for privilege escalation from multiple users
            privilege_escalations = [
                e for e in window_events
                if e.is_privilege_escalation
            ]
            
            if len(privilege_escalations) > 1:
                unique_users = set(e.user for e in privilege_escalations if e.user)
                if len(unique_users) > 1:
                    correlations.append(Correlation(
                        events=privilege_escalations,
                        correlation_type="coordinated_privilege_escalation",
                        description=f"Multiple users ({len(unique_users)}) escalated privileges within {self.CORRELATION_TIME_WINDOW}s",
                        risk_level=RiskLevel.HIGH,
                        timestamp=window_events[0].timestamp,
                        users_involved=unique_users,
                        auids_involved=set(e.auid for e in privilege_escalations)
                    ))
            
            # Check for destructive actions from multiple users
            destructive_actions = [
                e for e in window_events
                if e.is_destructive
            ]
            
            if len(destructive_actions) > 1:
                unique_users = set(e.user for e in destructive_actions if e.user)
                if len(unique_users) > 1:
                    correlations.append(Correlation(
                        events=destructive_actions,
                        correlation_type="coordinated_destructive_action",
                        description=f"Multiple users ({len(unique_users)}) performed destructive actions within {self.CORRELATION_TIME_WINDOW}s",
                        risk_level=RiskLevel.HIGH,
                        timestamp=window_events[0].timestamp,
                        users_involved=unique_users,
                        auids_involved=set(e.auid for e in destructive_actions)
                    ))
        
        return correlations
    
    def _group_by_time_windows(self, events: List[AuditEvent]) -> List[List[AuditEvent]]:
        """Group events into time windows.
        
        Args:
            events: List of events
            
        Returns:
            List of event groups within time windows
        """
        if not events:
            return []
        
        sorted_events = sorted(events, key=lambda e: e.timestamp)
        windows: List[List[AuditEvent]] = []
        current_window: List[AuditEvent] = [sorted_events[0]]
        window_start = sorted_events[0].timestamp
        
        for event in sorted_events[1:]:
            time_diff = (event.timestamp - window_start).total_seconds()
            
            if time_diff <= self.CORRELATION_TIME_WINDOW:
                current_window.append(event)
            else:
                if len(current_window) > 1:
                    windows.append(current_window)
                current_window = [event]
                window_start = event.timestamp
        
        if len(current_window) > 1:
            windows.append(current_window)
        
        return windows
    
    def _find_similar_commands(self, events: List[AuditEvent]) -> List[Correlation]:
        """Find similar commands executed by different users.
        
        Args:
            events: Events in a time window
            
        Returns:
            List of correlations
        """
        correlations: List[Correlation] = []
        
        # Group by command pattern
        command_groups: Dict[str, List[AuditEvent]] = defaultdict(list)
        
        for event in events:
            if not event.command:
                continue
            
            # Normalize command (remove arguments, keep base command)
            base_command = self._extract_base_command(event.command)
            if base_command:
                command_groups[base_command].append(event)
        
        # Find commands executed by multiple users
        for base_cmd, cmd_events in command_groups.items():
            unique_users = set(e.user for e in cmd_events if e.user)
            
            if len(unique_users) > 1 and len(cmd_events) >= 2:
                # Check if it's a suspicious command
                if self._is_suspicious_command(base_cmd):
                    correlations.append(Correlation(
                        events=cmd_events,
                        correlation_type="coordinated_command",
                        description=f"Multiple users ({len(unique_users)}) executed similar command: {base_cmd}",
                        risk_level=RiskLevel.MEDIUM,
                        timestamp=cmd_events[0].timestamp,
                        users_involved=unique_users,
                        auids_involved=set(e.auid for e in cmd_events)
                    ))
        
        return correlations
    
    def _extract_base_command(self, command: str) -> Optional[str]:
        """Extract base command from full command string.
        
        Args:
            command: Full command string
            
        Returns:
            Base command name or None
        """
        if not command:
            return None
        
        # Split and get first part
        parts = command.strip().split()
        if parts:
            # Remove path, keep just command name
            base = parts[0].split('/')[-1]
            return base.lower()
        
        return None
    
    def _is_suspicious_command(self, command: str) -> bool:
        """Check if command is suspicious for coordination.
        
        Args:
            command: Base command name
            
        Returns:
            True if suspicious
        """
        suspicious_commands = [
            'curl', 'wget', 'nc', 'netcat', 'ssh', 'scp',
            'rm', 'chmod', 'chown', 'mv', 'cp',
            'python', 'perl', 'ruby', 'bash', 'sh'
        ]
        
        return command.lower() in suspicious_commands
    
    def detect_lateral_movement(self, events: List[AuditEvent]) -> List[LateralMovement]:
        """Detect lateral movement patterns.
        
        Args:
            events: List of audit events
            
        Returns:
            List of lateral movement detections
        """
        movements: List[LateralMovement] = []
        
        # Group events by user
        events_by_user: Dict[Optional[int], List[AuditEvent]] = defaultdict(list)
        for event in events:
            events_by_user[event.auid].append(event)
        
        # Look for su/ssh patterns indicating user switching
        from models import EventType
        for auid, user_events in events_by_user.items():
            for event in user_events:
                if event.event_type == EventType.SU and event.command:
                    # Try to extract target user
                    target_user = self._extract_target_user(event.command)
                    if target_user and target_user != event.user:
                        movements.append(LateralMovement(
                            source_user=event.user,
                            target_user=target_user,
                            source_auid=auid,
                            target_auid=None,  # Would need to resolve
                            events=[event],
                            timestamp=event.timestamp,
                            risk_level=RiskLevel.MEDIUM
                        ))
        
        return movements
    
    def _extract_target_user(self, command: str) -> Optional[str]:
        """Extract target user from su/ssh command.
        
        Args:
            command: Command string
            
        Returns:
            Target username or None
        """
        # Pattern: su - username, su username, ssh user@host
        patterns = [
            r'su\s+-\s+(\w+)',
            r'su\s+(\w+)',
            r'ssh\s+(\w+)@',
        ]
        
        import re
        for pattern in patterns:
            match = re.search(pattern, command)
            if match:
                return match.group(1)
        
        return None
    
    def correlate_by_ip(self, events: List[AuditEvent]) -> List[Correlation]:
        """Correlate events by IP address.
        
        Args:
            events: List of audit events
            
        Returns:
            List of IP-based correlations
        """
        correlations: List[Correlation] = []
        
        # Group by IP
        ip_groups: Dict[str, List[AuditEvent]] = defaultdict(list)
        
        for event in events:
            if event.dst_ip:
                ip_groups[event.dst_ip].append(event)
            elif event.src_ip:
                ip_groups[event.src_ip].append(event)
        
        # Find IPs accessed by multiple users
        for ip, ip_events in ip_groups.items():
            unique_users = set(e.user for e in ip_events if e.user)
            
            if len(unique_users) > 1:
                correlations.append(Correlation(
                    events=ip_events,
                    correlation_type="shared_ip_access",
                    description=f"Multiple users ({len(unique_users)}) accessed IP: {ip}",
                    risk_level=RiskLevel.MEDIUM,
                    timestamp=ip_events[0].timestamp,
                    users_involved=unique_users,
                    auids_involved=set(e.auid for e in ip_events)
                ))
        
        return correlations