#!/bin/bash
# Test script to verify monitor mode works with tshark
# Usage: ./test_monitor_tshark.sh [interface] [channel] [duration_seconds] [--keep-pcap]
# Or: KEEP_PCAP=1 sudo -E ./test_monitor_tshark.sh [interface] [channel] [duration_seconds]

set -e

# Parse arguments
KEEP_PCAP_FLAG=""
ARGS=()
for arg in "$@"; do
    if [ "$arg" = "--keep-pcap" ] || [ "$arg" = "-k" ]; then
        KEEP_PCAP_FLAG="1"
    else
        ARGS+=("$arg")
    fi
done

INTERFACE="${ARGS[0]:-wlan0}"
CHANNEL="${ARGS[1]:-36}"
DURATION="${ARGS[2]:-10}"  # Default 10 seconds, minimum 1 second

# Check for KEEP_PCAP environment variable or flag
if [ -n "$KEEP_PCAP_FLAG" ]; then
    KEEP_PCAP="1"
elif [ -z "${KEEP_PCAP:-}" ]; then
    KEEP_PCAP=""
fi

# Ensure minimum 1 second
if [ "$DURATION" -lt 1 ]; then
    DURATION=1
fi

echo "=== Testing Monitor Mode with tshark ==="
echo "Interface: $INTERFACE"
echo "Channel: $CHANNEL"
echo "Duration: $DURATION seconds"
echo ""

# Check if running as root
if [ "$EUID" -ne 0 ]; then 
    echo "Please run as root (use sudo)"
    exit 1
fi

# Check if tshark is installed
if ! command -v tshark &> /dev/null; then
    echo "tshark is not installed. Installing..."
    if command -v apt-get &> /dev/null; then
        sudo apt-get update
        sudo apt-get install -y tshark
    elif command -v dnf &> /dev/null; then
        sudo dnf install -y wireshark-cli
    else
        echo "Please install tshark manually"
        exit 1
    fi
fi

# Unmanage from NetworkManager
if command -v nmcli &> /dev/null; then
    echo "Unmanaging interface from NetworkManager..."
    nmcli device set "$INTERFACE" managed no 2>/dev/null || true
fi

# Unblock WiFi
rfkill unblock wifi 2>/dev/null || true

# Bring down interface
echo "Bringing down interface..."
ip link set "$INTERFACE" down 2>/dev/null || true
sleep 0.5

# Set monitor mode
echo "Setting monitor mode..."
if ! iw dev "$INTERFACE" set type monitor; then
    echo "Error: Failed to set monitor mode"
    exit 1
fi

sleep 0.5

# Bring up interface
echo "Bringing up interface..."
ip link set "$INTERFACE" up || echo "Warning: Failed to bring interface up"

sleep 0.5

# Set channel
echo "Setting channel to $CHANNEL..."
iw dev "$INTERFACE" set channel "$CHANNEL" || echo "Warning: Failed to set channel"

# Verify monitor mode
echo ""
echo "Verifying monitor mode..."
iw dev "$INTERFACE" info | grep -E "(type|channel)" || echo "Could not verify"

# Check DLT with tshark (capture for 1 second)
echo ""
echo "Checking Data Link Type (1 second test capture)..."
echo "(This may take up to 2 seconds if no packets are present)"

# Use timeout with -c to limit packets and avoid hanging
# Capture both stdout and stderr
# Fields: frame.number, wlan.ra (Receiver Address), wlan.ta (Transmitter Address), radiotap.present
# Note: wlan.ra and wlan.ta may not be available for all frame types, but are correct for monitor mode
TEST_OUTPUT=$(timeout 2 tshark -i "$INTERFACE" -c 100 -T fields \
    -e frame.number \
    -e wlan.ra \
    -e wlan.ta \
    -e radiotap.present \
    2>&1 || true)
TEST_EXIT_CODE=${PIPESTATUS[0]}

# Show any warnings/errors from tshark (but not packet data)
echo "$TEST_OUTPUT" | grep -E "(Running as|tshark:|Warning|Error|Capturing)" || true

# Count packets (lines starting with a number, excluding error messages)
PACKET_LINES=$(echo "$TEST_OUTPUT" | grep -E '^[0-9]+' || true)
PACKET_COUNT=$(echo "$PACKET_LINES" | wc -l || echo "0")

# Count lines with radiotap.present (4th field) set to 1
PLCP_COUNT=$(echo "$PACKET_LINES" | awk -F'\t' 'NF >= 4 && $1 != "" && $4 != "" && $4 != "0" && $4 != "-" {count++} END {print count+0}' || echo "0")

# Show sample output with better formatting
if [ "$PACKET_COUNT" -gt 0 ]; then
    echo "Sample packets:"
    echo "$PACKET_LINES" | head -5 | awk -F'\t' '{
        ra = ($2 != "" && $2 != "-") ? $2 : "N/A"
        ta = ($3 != "" && $3 != "-") ? $3 : "N/A"
        radiotap = ($4 == "1" || $4 == "1.0") ? "yes" : "no"
        printf "  Frame %s: RA=%s, TA=%s, PLCP=%s\n", $1, ra, ta, radiotap
    }'
fi

echo ""
echo "Test capture results:"
echo "  Packets captured: $PACKET_COUNT"
echo "  PLCP headers: $PLCP_COUNT"
if [ "$PLCP_COUNT" -eq 0 ] && [ "$PACKET_COUNT" -gt 0 ]; then
    echo "  Note: Packets captured but no radiotap headers (may be using DLT_IEEE802_11 instead of DLT_IEEE802_11_RADIO)"
fi

echo ""
echo "=== Starting tshark capture ($DURATION seconds) ==="
echo "Press Ctrl+C to stop early"
echo ""

# Capture for specified duration and count packets
echo "Capturing packets for $DURATION seconds..."

# Use a temporary pcap file to avoid field extraction errors during capture
# Capture to file first, then parse it - this prevents tshark from exiting early
TEMP_PCAP=$(mktemp /tmp/tshark_capture_XXXXXX.pcap)
echo "Capturing to temporary file: $TEMP_PCAP"

set +e  # Temporarily disable exit on error
# Capture to pcap file - this won't error on missing fields
# Use -b filesize:100000 to rotate files if needed, but we'll only use the first
timeout "$DURATION" tshark -q -i "$INTERFACE" -n -w "$TEMP_PCAP" 2>/dev/null
CAPTURE_EXIT_CODE=$?
set -e  # Re-enable exit on error

# Force sync to ensure file is written
sync

# Exit code 124 means timeout occurred (expected), 0 means command completed normally
if [ "$CAPTURE_EXIT_CODE" -ne 0 ] && [ "$CAPTURE_EXIT_CODE" -ne 124 ]; then
    echo "Warning: tshark capture exited with code $CAPTURE_EXIT_CODE"
fi

# Check if pcap file exists and get its size
if [ -f "$TEMP_PCAP" ]; then
    PCAP_SIZE=$(stat -c%s "$TEMP_PCAP" 2>/dev/null || stat -f%z "$TEMP_PCAP" 2>/dev/null || echo "0")
    echo "Pcap file size: $PCAP_SIZE bytes"
    
    # Count packets in raw pcap file using capinfos or tshark
    if command -v capinfos &> /dev/null; then
        RAW_PACKET_COUNT=$(capinfos -c "$TEMP_PCAP" 2>/dev/null | grep "^Number of packets:" | awk '{print $4}' || echo "0")
        echo "Raw packets in pcap file: $RAW_PACKET_COUNT"
    else
        # Fallback: use tshark to count packets
        RAW_PACKET_COUNT=$(tshark -q -r "$TEMP_PCAP" -n -T fields -e frame.number 2>/dev/null | tail -1 || echo "0")
        if [ "$RAW_PACKET_COUNT" != "0" ] && [ -n "$RAW_PACKET_COUNT" ]; then
            echo "Raw packets in pcap file: $RAW_PACKET_COUNT"
        else
            echo "Raw packets in pcap file: (could not determine)"
        fi
    fi
else
    echo "ERROR: Pcap file was not created: $TEMP_PCAP"
    exit 1
fi

# Now parse the pcap file to extract fields
# Only extract 802.11 header fields - data payloads are encrypted
# Include PHY rate and MCS for histograms
# Note: Some fields may not exist for all frames - tshark will output empty values
CAPTURE_OUTPUT=$(tshark -q -r "$TEMP_PCAP" -n -T fields \
    -E occurrence=f \
    -E separator='\t' \
    -E header=n \
    -e frame.number \
    -e frame.time \
    -e wlan.ra \
    -e wlan.ta \
    -e wlan.fc.type \
    -e wlan.fc.subtype \
    -e wlan.fc.type_subtype \
    -e wlan.fc.protected \
    -e wlan.fc.retry \
    -e wlan.duration \
    -e radiotap.present \
    -e radiotap.datarate \
    -e radiotap.mcs.index \
    -e wlan_radio.data_rate \
    -e wlan_radio.mcs.index \
    2>&1 | grep -E '^[0-9]' | head -1000000 || true)

# Clean up temp file (unless KEEP_PCAP is set)
if [ -z "$KEEP_PCAP" ]; then
    echo "Cleaning up temporary pcap file: $TEMP_PCAP"
    rm -f "$TEMP_PCAP"
else
    echo "Keeping temporary pcap file: $TEMP_PCAP"
    echo "  (Use: tshark -r $TEMP_PCAP to analyze)"
fi

# Force output flush
sync

# Show stats immediately after capture completes
echo ""
echo "=== Capture Statistics ==="

# Show any warnings/errors (but not the "Running as root" or "Capturing" messages)
WARNINGS=$(echo "$CAPTURE_OUTPUT" | grep -E "(tshark:|Warning|Error)" | grep -v -E "(Running as|Capturing)" | head -5 || true)
if [ -n "$WARNINGS" ]; then
    echo "$WARNINGS"
    echo ""
fi

# Count total packets captured (lines starting with a number followed by tab)
# Filter out tshark status messages like "100 packets captured" or "Capturing on..."
# Only count lines that look like actual packet data: number, tab, then more fields
# Also handle lines that start with just a number (frame.number field)
PACKET_LINES=$(echo "$CAPTURE_OUTPUT" | grep -E '^[0-9]+' | grep -v -E '(packets captured|Capturing on|Running as|tshark:)' || true)
FINAL_COUNT=$(echo "$PACKET_LINES" | wc -l || echo "0")

# If we got very few packets but raw count shows many, there might be a parsing issue
if [ "$FINAL_COUNT" -lt "$RAW_PACKET_COUNT" ] && [ "$RAW_PACKET_COUNT" -gt 10 ]; then
    echo "Warning: Parsed $FINAL_COUNT packets but pcap file contains $RAW_PACKET_COUNT packets"
    echo "  This may indicate field extraction issues. Check tshark output above."
fi

# Count packets with PLCP headers (radiotap present)
# radiotap.present is field 11 (after frame.number, frame.time, wlan.ra, wlan.ta, wlan.fc.type, wlan.fc.subtype, wlan.fc.type_subtype, wlan.fc.protected, wlan.fc.retry, wlan.duration)
PLCP_COUNT=$(echo "$PACKET_LINES" | awk -F'\t' 'NF >= 11 && $1 != "" && $11 != "" && $11 != "0" && $11 != "-" {count++} END {print count+0}' || echo "0")

# Display stats immediately - always show these
echo "Total packets captured: $FINAL_COUNT"
echo "PLCP headers: $PLCP_COUNT"
if [ "$FINAL_COUNT" -gt 0 ]; then
    # Calculate rate
    RATE=$(awk "BEGIN {printf \"%.1f\", $FINAL_COUNT / $DURATION}")
    echo "Packet rate: $RATE packets/second"
fi
echo ""

# Display sample packets with readable format
if [ -n "$PACKET_LINES" ] && [ "$FINAL_COUNT" -gt 0 ]; then
    echo "Sample packets (first 10):"
    echo "$PACKET_LINES" | head -10 | awk -F'\t' '{
        ra = ($3 != "" && $3 != "-") ? $3 : "N/A"
        ta = ($4 != "" && $4 != "-") ? $4 : "N/A"
        type = ($5 != "" && $5 != "-") ? $5 : "N/A"
        subtype = ($6 != "" && $6 != "-") ? $6 : "N/A"
        protected = ($8 == "1" || $8 == "1.0") ? "encrypted" : "unencrypted"
        retry = ($9 == "1" || $9 == "1.0") ? "retry" : ""
        duration = ($10 != "" && $10 != "-") ? $10 : "N/A"
        radiotap = ($11 == "1" || $11 == "1.0") ? "yes" : (($11 != "" && $11 != "-") ? "no" : "N/A")
        retry_str = (retry != "") ? sprintf(" [%s]", retry) : ""
        printf "  Frame %s: RA=%s, TA=%s, type=%s/%s, %s, dur=%s, PLCP=%s%s\n",
            $1, ra, ta, type, subtype, protected, duration, radiotap, retry_str
    }'
    echo ""
    
    # Count unique RA/TA pairs
    echo "Unique RA/TA pairs (with counts):"
    UNIQUE_PAIRS=$(echo "$PACKET_LINES" | awk -F'\t' '{
        ra = ($3 != "" && $3 != "-") ? $3 : "N/A"
        ta = ($4 != "" && $4 != "-") ? $4 : "N/A"
        if (ra != "N/A" || ta != "N/A") {
            pair = ra " -> " ta
            count[pair]++
        }
    }
    END {
        for (pair in count) {
            printf "%d\t%s\n", count[pair], pair
        }
    }' | sort -rn)
    
    if [ -n "$UNIQUE_PAIRS" ]; then
        echo "$UNIQUE_PAIRS" | awk -F'\t' '{
            printf "  %s: %d frame(s)\n", $2, $1
        }'
    else
        echo "  (no valid RA/TA pairs found)"
    fi
    echo ""
    
    # Generate PHY rate and MCS histograms per RA/TA pair
    echo "PHY Rate and MCS Histograms per RA/TA pair:"
    echo "$PACKET_LINES" | awk -F'\t' '{
        ra = ($3 != "" && $3 != "-") ? $3 : "N/A"
        ta = ($4 != "" && $4 != "-") ? $4 : "N/A"
        if (ra != "N/A" || ta != "N/A") {
            pair = ra " -> " ta
            
            # Get PHY rate (try radiotap.datarate first, then wlan_radio.data_rate)
            phy_rate = ""
            if ($12 != "" && $12 != "-" && $12 != "0") {
                phy_rate = $12
            } else if ($14 != "" && $14 != "-" && $14 != "0") {
                phy_rate = $14
            }
            
            # Get MCS index (try radiotap.mcs.index first, then wlan_radio.mcs.index)
            mcs = ""
            if ($13 != "" && $13 != "-" && $13 != "0" && $13 != "255") {
                mcs = $13
            } else if ($15 != "" && $15 != "-" && $15 != "0" && $15 != "255") {
                mcs = $15
            }
            
            # Only count data frames (type 2) for histograms
            type = ($5 != "" && $5 != "-") ? $5 : ""
            if (type == "2") {
                if (phy_rate != "") {
                    rate_hist[pair][phy_rate]++
                }
                if (mcs != "") {
                    mcs_hist[pair][mcs]++
                }
            }
        }
    }
    END {
        # Process each pair
        for (pair in rate_hist) {
            printf "\n  %s:\n", pair
            
            # PHY Rate histogram
            printf "    PHY Rate (Mbps):\n"
            # Collect rates and sort
            n = 0
            for (rate in rate_hist[pair]) {
                rates[n++] = rate + 0
                rate_map[rate + 0] = rate
            }
            # Sort rates
            for (i = 0; i < n; i++) {
                for (j = i + 1; j < n; j++) {
                    if (rates[i] > rates[j]) {
                        tmp = rates[i]
                        rates[i] = rates[j]
                        rates[j] = tmp
                    }
                }
            }
            for (i = 0; i < n; i++) {
                rate_val = rates[i]
                rate_str = rate_map[rate_val]
                printf "      %s Mbps: %d frame(s)\n", rate_str, rate_hist[pair][rate_str]
            }
            
            # MCS histogram
            if (pair in mcs_hist) {
                printf "    MCS Index:\n"
                m = 0
                for (mcs in mcs_hist[pair]) {
                    mcs_vals[m++] = mcs + 0
                    mcs_map[mcs + 0] = mcs
                }
                # Sort MCS
                for (i = 0; i < m; i++) {
                    for (j = i + 1; j < m; j++) {
                        if (mcs_vals[i] > mcs_vals[j]) {
                            tmp = mcs_vals[i]
                            mcs_vals[i] = mcs_vals[j]
                            mcs_vals[j] = tmp
                        }
                    }
                }
                for (i = 0; i < m; i++) {
                    mcs_val = mcs_vals[i]
                    mcs_str = mcs_map[mcs_val]
                    printf "      MCS %s: %d frame(s)\n", mcs_str, mcs_hist[pair][mcs_str]
                }
            } else {
                printf "    MCS Index: (no MCS data)\n"
            }
        }
    }'
    echo ""
    
    # Frame type breakdown
    echo "Frame type breakdown:"
    echo "$PACKET_LINES" | awk -F'\t' '{
        type = ($5 != "" && $5 != "-") ? $5 : "unknown"
        subtype = ($6 != "" && $6 != "-") ? $6 : "unknown"
        type_name = "Unknown"
        if (type == "0") type_name = "Management"
        else if (type == "1") type_name = "Control"
        else if (type == "2") type_name = "Data"
        count[type_name]++
    }
    END {
        for (t in count) {
            printf "  %s: %d frame(s)\n", t, count[t]
        }
    }' | sort -rn
    echo ""
    
    # Analyze data frames (iperf uses QoS Data frames, subtype 8)
    echo "Data frame analysis (iperf typically uses QoS Data frames, subtype 8):"
    DATA_FRAMES=$(echo "$PACKET_LINES" | awk -F'\t' '{
        type = ($5 != "" && $5 != "-") ? $5 : ""
        subtype = ($6 != "" && $6 != "-") ? $6 : ""
        if (type == "2" && subtype == "8") {  # QoS Data frames
            print $0
        }
    }')
    DATA_COUNT=$(echo "$DATA_FRAMES" | wc -l || echo "0")
    echo "  QoS Data frames (type 2, subtype 8): $DATA_COUNT"
    
    # Count encrypted vs unencrypted data frames
    ENCRYPTED_DATA=$(echo "$DATA_FRAMES" | awk -F'\t' '$8 == "1" || $8 == "1.0" {count++} END {print count+0}')
    UNENCRYPTED_DATA=$(echo "$DATA_FRAMES" | awk -F'\t' '$8 != "1" && $8 != "1.0" && $8 != "" && $8 != "-" {count++} END {print count+0}')
    echo "    Encrypted: $ENCRYPTED_DATA"
    echo "    Unencrypted: $UNENCRYPTED_DATA"
    
    if [ "$DATA_COUNT" -gt 0 ]; then
        echo "  Sample QoS Data frames (likely iperf traffic):"
        echo "$DATA_FRAMES" | head -5 | awk -F'\t' '{
            ra = ($3 != "" && $3 != "-") ? $3 : "N/A"
            ta = ($4 != "" && $4 != "-") ? $4 : "N/A"
            protected = ($8 == "1" || $8 == "1.0") ? "encrypted" : "unencrypted"
            retry = ($9 == "1" || $9 == "1.0") ? "retry" : ""
            duration = ($10 != "" && $10 != "-") ? $10 : "N/A"
            retry_str = (retry != "") ? sprintf(" [%s]", retry) : ""
            printf "    Frame %s: RA=%s, TA=%s, %s, dur=%s%s\n",
                $1, ra, ta, protected, duration, retry_str
        }'
    fi
    echo ""
    
    # Frames involving server MAC (80:84:89:93:c4:b6)
    echo "Frames involving server MAC (80:84:89:93:c4:b6):"
    SERVER_MAC="80:84:89:93:c4:b6"
    SERVER_FRAMES=$(echo "$PACKET_LINES" | awk -F'\t' -v mac="$SERVER_MAC" '{
        ra = ($3 != "" && $3 != "-") ? $3 : ""
        ta = ($4 != "" && $4 != "-") ? $4 : ""
        if (ra == mac || ta == mac) {
            print $0
        }
    }')
    SERVER_COUNT=$(echo "$SERVER_FRAMES" | wc -l || echo "0")
    echo "  Total frames with server MAC: $SERVER_COUNT"
    if [ "$SERVER_COUNT" -gt 0 ]; then
        echo "  Frame type breakdown:"
        echo "$SERVER_FRAMES" | awk -F'\t' '{
            type = ($5 != "" && $5 != "-") ? $5 : "unknown"
            subtype = ($6 != "" && $6 != "-") ? $6 : "unknown"
            type_name = "Unknown"
            if (type == "0") type_name = "Management"
            else if (type == "1") type_name = "Control"
            else if (type == "2") type_name = "Data"
            count[type_name]++
        }
        END {
            for (t in count) {
                printf "    %s: %d frame(s)\n", t, count[t]
            }
        }' | sort -rn
        echo "  Sample frames:"
        echo "$SERVER_FRAMES" | head -5 | awk -F'\t' '{
            ra = ($3 != "" && $3 != "-") ? $3 : "N/A"
            ta = ($4 != "" && $4 != "-") ? $4 : "N/A"
            type = ($5 != "" && $5 != "-") ? $5 : "N/A"
            subtype = ($6 != "" && $6 != "-") ? $6 : "N/A"
            protected = ($8 == "1" || $8 == "1.0") ? "encrypted" : "unencrypted"
            retry = ($9 == "1" || $9 == "1.0") ? "retry" : ""
            duration = ($10 != "" && $10 != "-") ? $10 : "N/A"
            retry_str = (retry != "") ? sprintf(" [%s]", retry) : ""
            printf "    Frame %s: RA=%s, TA=%s, type=%s/%s, %s, dur=%s%s\n",
                $1, ra, ta, type, subtype, protected, duration, retry_str
        }'
    fi
    echo ""
else
    echo "(No packets captured)"
    echo ""
fi

echo "=== Summary ==="
if [ "$FINAL_COUNT" -gt 0 ]; then
    echo "✓ Monitor mode is working! Captured $FINAL_COUNT packet(s)"
    if [ "$PLCP_COUNT" -gt 0 ]; then
        echo "✓ PLCP headers detected: $PLCP_COUNT packet(s) with radiotap information"
    else
        echo "⚠ No PLCP headers detected (may be using DLT_IEEE802_11 instead of DLT_IEEE802_11_RADIO)"
    fi
else
    echo "✗ No packets captured. Check:"
    echo "  1. Is there WiFi traffic on channel $CHANNEL?"
    echo "  2. Is the interface actually in monitor mode? (iw dev $INTERFACE info)"
    echo "  3. Try a different channel (e.g., 1, 6, 11 for 2.4GHz)"
    echo "  4. Try a longer duration: sudo ./test_monitor_tshark.sh $INTERFACE $CHANNEL 30"
fi