#!/bin/bash
#
# capture_and_plot.sh - Capture and visualize CVE-2023-1206 trigger traffic
#
# This script:
# 1. Captures network traffic on the Docker bridge
# 2. Extracts packet timing data
# 3. Generates visualizations of the burst patterns
#

set -e

# Auto-detect Docker bridge for triple_cve_net
detect_interface() {
    # First try to get the bridge name from the network
    local bridge=$(docker network inspect triple_cve_net 2>/dev/null | grep -oP '"com.docker.network.bridge.name": "\K[^"]+' || echo "")
    
    if [ -z "$bridge" ]; then
        # Try to find br-* interface from network ID
        local net_id=$(docker network inspect triple_cve_net 2>/dev/null | grep -oP '"Id": "\K[^"]+' | head -c 12 || echo "")
        if [ -n "$net_id" ]; then
            bridge="br-$net_id"
        fi
    fi
    
    if [ -z "$bridge" ]; then
        # Fallback: find any br-* interface
        bridge=$(ip link show 2>/dev/null | grep -oP 'br-[a-f0-9]+' | head -1)
    fi
    
    if [ -z "$bridge" ]; then
        # Last resort
        bridge="docker0"
    fi
    
    echo "$bridge"
}

INTERFACE="${1:-$(detect_interface)}"
CAPTURE_TIME="${2:-30}"
OUTPUT_DIR="./captures"
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
PCAP_FILE="${OUTPUT_DIR}/triple_cve_${TIMESTAMP}.pcap"
CSV_FILE="${OUTPUT_DIR}/triple_cve_${TIMESTAMP}.csv"
PLOT_FILE="${OUTPUT_DIR}/triple_cve_${TIMESTAMP}.png"

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
NC='\033[0m'

log_info() { echo -e "${GREEN}[+]${NC} $1"; }
log_warn() { echo -e "${YELLOW}[!]${NC} $1"; }
log_error() { echo -e "${RED}[-]${NC} $1"; }

mkdir -p "$OUTPUT_DIR"

cat << 'BANNER'
╔══════════════════════════════════════════════════════════════╗
║  CVE-2023-1206 Traffic Capture & Visualization               ║
║  Triple CVE Covert Channel Analysis                          ║
╚══════════════════════════════════════════════════════════════╝

BANNER

# Check dependencies
check_deps() {
    local missing=()
    
    command -v tshark >/dev/null 2>&1 || missing+=("tshark (wireshark-cli)")
    command -v python3 >/dev/null 2>&1 || missing+=("python3")
    
    if [ ${#missing[@]} -gt 0 ]; then
        log_error "Missing dependencies: ${missing[*]}"
        echo "Install with: sudo apt install wireshark-cli python3-matplotlib python3-pandas"
        exit 1
    fi
    
    # Check Python packages
    python3 -c "import matplotlib" 2>/dev/null || {
        log_warn "matplotlib not found, installing..."
        pip3 install matplotlib --break-system-packages 2>/dev/null || pip3 install matplotlib
    }
    
    python3 -c "import pandas" 2>/dev/null || {
        log_warn "pandas not found, installing..."
        pip3 install pandas --break-system-packages 2>/dev/null || pip3 install pandas
    }
}

# Capture traffic
capture_traffic() {
    log_info "Capturing on interface: $INTERFACE for ${CAPTURE_TIME}s"
    log_info "Output: $PCAP_FILE"
    echo ""
    log_warn "Start the triple CVE demo now!"
    echo "  Terminal 1: ./run_triple_cve.sh responder"
    echo "  Terminal 2: sudo ./run_triple_cve.sh initiator 'SECRET MESSAGE'"
    echo ""
    
    # Capture UDP traffic on our ports
    sudo tshark -i "$INTERFACE" -a duration:"$CAPTURE_TIME" \
        -f "udp port 31336 or udp port 31337" \
        -w "$PCAP_FILE" 2>/dev/null &
    
    TSHARK_PID=$!
    
    # Show countdown
    for i in $(seq $CAPTURE_TIME -1 1); do
        printf "\r${CYAN}[*]${NC} Capturing... %3d seconds remaining " $i
        sleep 1
    done
    echo ""
    
    wait $TSHARK_PID 2>/dev/null || true
    
    log_info "Capture complete: $PCAP_FILE"
}

# Extract data to CSV
extract_csv() {
    log_info "Extracting packet data to CSV..."
    
    tshark -r "$PCAP_FILE" -T fields \
        -e frame.time_relative \
        -e frame.time_delta \
        -e ip.src \
        -e ip.dst \
        -e udp.srcport \
        -e udp.dstport \
        -e udp.length \
        -e data.data \
        -E header=y -E separator=, \
        > "$CSV_FILE" 2>/dev/null
    
    PACKET_COUNT=$(wc -l < "$CSV_FILE")
    log_info "Extracted $((PACKET_COUNT - 1)) packets to $CSV_FILE"
}

# Generate visualization
generate_plot() {
    log_info "Generating visualization..."
    
    python3 << PYTHON_SCRIPT
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import numpy as np
from datetime import datetime
import sys

# Read CSV
try:
    df = pd.read_csv("$CSV_FILE")
    if len(df) == 0:
        print("No packets captured!")
        sys.exit(1)
except Exception as e:
    print(f"Error reading CSV: {e}")
    sys.exit(1)

# Rename columns for easier access
df.columns = ['time_rel', 'time_delta', 'src_ip', 'dst_ip', 'src_port', 'dst_port', 'length', 'data']

# Convert time to float
df['time_rel'] = pd.to_numeric(df['time_rel'], errors='coerce')
df['time_delta'] = pd.to_numeric(df['time_delta'], errors='coerce')

# Identify packet types by magic numbers in data
def classify_packet(row):
    data = str(row.get('data', ''))
    port = row.get('dst_port', 0)
    
    # CVE-2023-1206 trigger magic: 0xCAFE1206 (little endian: 0612feca)
    if '0612feca' in data.lower() or 'cafe1206' in data.lower():
        return 'TRIGGER (CVE-2023-1206)'
    # Sync magic: 0xDEAD0040 (little endian: 4000adde)
    elif '4000adde' in data.lower() or 'dead0040' in data.lower():
        return 'SYNC (Key Agreement)'
    elif port == 31336:
        return 'TRIGGER (CVE-2023-1206)'
    elif port == 31337:
        return 'SYNC (Key Agreement)'
    else:
        return 'OTHER'

df['packet_type'] = df.apply(classify_packet, axis=1)

# Create figure with multiple subplots
fig = plt.figure(figsize=(16, 12))
fig.suptitle('Triple CVE Covert Channel - Traffic Analysis\\nCVE-2023-1206 (Trigger) + CVE-2025-40040 (KSM Key Agreement)', 
             fontsize=14, fontweight='bold')

# Color scheme
colors = {
    'TRIGGER (CVE-2023-1206)': '#FF4444',
    'SYNC (Key Agreement)': '#4444FF', 
    'OTHER': '#888888'
}

# Subplot 1: Packet timeline (scatter plot)
ax1 = fig.add_subplot(3, 2, 1)
for ptype in df['packet_type'].unique():
    mask = df['packet_type'] == ptype
    ax1.scatter(df[mask]['time_rel'], [1]*sum(mask), 
                c=colors.get(ptype, '#888888'), 
                label=ptype, alpha=0.6, s=20)
ax1.set_xlabel('Time (seconds)')
ax1.set_ylabel('Packets')
ax1.set_title('Packet Timeline')
ax1.legend(loc='upper right', fontsize=8)
ax1.set_yticks([])
ax1.grid(True, alpha=0.3)

# Subplot 2: Packets per 100ms (histogram showing bursts)
ax2 = fig.add_subplot(3, 2, 2)
max_time = df['time_rel'].max()
bins = np.arange(0, max_time + 0.1, 0.1)  # 100ms bins

for ptype in ['TRIGGER (CVE-2023-1206)', 'SYNC (Key Agreement)']:
    mask = df['packet_type'] == ptype
    if sum(mask) > 0:
        ax2.hist(df[mask]['time_rel'], bins=bins, alpha=0.7, 
                 label=ptype, color=colors.get(ptype, '#888888'))

ax2.set_xlabel('Time (seconds)')
ax2.set_ylabel('Packets per 100ms')
ax2.set_title('Packet Rate Over Time (100ms bins) - Shows CVE-2023-1206 Burst')
ax2.legend(loc='upper right', fontsize=8)
ax2.grid(True, alpha=0.3)

# Subplot 3: Inter-packet timing (to show burst pattern)
ax3 = fig.add_subplot(3, 2, 3)
trigger_mask = df['packet_type'] == 'TRIGGER (CVE-2023-1206)'
if sum(trigger_mask) > 1:
    trigger_times = df[trigger_mask]['time_rel'].values
    inter_arrival = np.diff(trigger_times) * 1000  # Convert to ms
    ax3.plot(range(len(inter_arrival)), inter_arrival, 'r-', alpha=0.7, linewidth=0.5)
    ax3.scatter(range(len(inter_arrival)), inter_arrival, c='red', s=5, alpha=0.5)
    ax3.set_xlabel('Packet Number')
    ax3.set_ylabel('Inter-arrival Time (ms)')
    ax3.set_title('CVE-2023-1206 Trigger - Inter-Packet Timing')
    ax3.set_ylim(0, max(inter_arrival.max() * 1.1, 1))
else:
    ax3.text(0.5, 0.5, 'No trigger packets captured', ha='center', va='center')
ax3.grid(True, alpha=0.3)

# Subplot 4: Sync packet timing (key agreement rounds)
ax4 = fig.add_subplot(3, 2, 4)
sync_mask = df['packet_type'] == 'SYNC (Key Agreement)'
if sum(sync_mask) > 1:
    sync_df = df[sync_mask].copy()
    sync_times = sync_df['time_rel'].values
    inter_arrival = np.diff(sync_times) * 1000
    ax4.plot(range(len(inter_arrival)), inter_arrival, 'b-', alpha=0.7, linewidth=0.5)
    ax4.scatter(range(len(inter_arrival)), inter_arrival, c='blue', s=5, alpha=0.5)
    ax4.set_xlabel('Packet Number')
    ax4.set_ylabel('Inter-arrival Time (ms)')
    ax4.set_title('Key Agreement Sync - Inter-Packet Timing (256 rounds)')
else:
    ax4.text(0.5, 0.5, 'No sync packets captured', ha='center', va='center')
ax4.grid(True, alpha=0.3)

# Subplot 5: Packet size distribution
ax5 = fig.add_subplot(3, 2, 5)
for ptype in df['packet_type'].unique():
    mask = df['packet_type'] == ptype
    if sum(mask) > 0:
        ax5.hist(df[mask]['length'].dropna(), bins=20, alpha=0.7,
                 label=ptype, color=colors.get(ptype, '#888888'))
ax5.set_xlabel('Packet Size (bytes)')
ax5.set_ylabel('Count')
ax5.set_title('Packet Size Distribution')
ax5.legend(loc='upper right', fontsize=8)
ax5.grid(True, alpha=0.3)

# Subplot 6: Statistics text
ax6 = fig.add_subplot(3, 2, 6)
ax6.axis('off')

stats_text = f"""
CAPTURE STATISTICS
══════════════════════════════════════

Total Packets: {len(df)}
Capture Duration: {df['time_rel'].max():.2f} seconds

PACKET BREAKDOWN
──────────────────────────────────────
"""

for ptype in df['packet_type'].unique():
    count = sum(df['packet_type'] == ptype)
    stats_text += f"  {ptype}: {count} packets\\n"

# Detect phases
trigger_packets = df[df['packet_type'] == 'TRIGGER (CVE-2023-1206)']
sync_packets = df[df['packet_type'] == 'SYNC (Key Agreement)']

if len(trigger_packets) > 0:
    trigger_start = trigger_packets['time_rel'].min()
    trigger_end = trigger_packets['time_rel'].max()
    trigger_duration = trigger_end - trigger_start
    trigger_rate = len(trigger_packets) / max(trigger_duration, 0.001)
    
    stats_text += f"""
PHASE 1: CVE-2023-1206 TRIGGER
──────────────────────────────────────
  Start: {trigger_start:.3f}s
  Duration: {trigger_duration*1000:.1f}ms
  Packet Rate: {trigger_rate:.0f} pkt/s
  Total: {len(trigger_packets)} packets
"""

if len(sync_packets) > 0:
    sync_start = sync_packets['time_rel'].min()
    sync_end = sync_packets['time_rel'].max()
    sync_duration = sync_end - sync_start
    
    stats_text += f"""
PHASE 2: KEY AGREEMENT (CVE-2025-40040)
──────────────────────────────────────
  Start: {sync_start:.3f}s
  Duration: {sync_duration:.1f}s
  Sync Packets: {len(sync_packets)}
  Expected: ~512 (256 rounds × 2 parties)
"""

ax6.text(0.05, 0.95, stats_text, transform=ax6.transAxes, 
         fontfamily='monospace', fontsize=9,
         verticalalignment='top',
         bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))

plt.tight_layout()
plt.savefig("$PLOT_FILE", dpi=150, bbox_inches='tight')
print(f"Plot saved to: $PLOT_FILE")

# Also save a summary
print("\\n" + "="*60)
print("CAPTURE SUMMARY")
print("="*60)
print(f"Total packets: {len(df)}")
for ptype in df['packet_type'].unique():
    print(f"  {ptype}: {sum(df['packet_type'] == ptype)}")
print("="*60)

PYTHON_SCRIPT
    
    log_info "Visualization saved to: $PLOT_FILE"
}

# Open plot
open_plot() {
    if command -v xdg-open >/dev/null 2>&1; then
        xdg-open "$PLOT_FILE" 2>/dev/null &
    elif command -v open >/dev/null 2>&1; then
        open "$PLOT_FILE" 2>/dev/null &
    fi
}

# Main
main() {
    echo "Interface: $INTERFACE"
    echo "Capture time: ${CAPTURE_TIME}s"
    echo "Output directory: $OUTPUT_DIR"
    echo ""
    
    check_deps
    capture_traffic
    extract_csv
    generate_plot
    open_plot
    
    echo ""
    log_info "Done! Files:"
    echo "  PCAP: $PCAP_FILE"
    echo "  CSV:  $CSV_FILE"
    echo "  Plot: $PLOT_FILE"
    echo ""
    log_info "Open in Wireshark: wireshark $PCAP_FILE"
}

main "$@"