#!/bin/bash
# Test script to verify monitor mode works with tshark
# Usage: ./test_monitor_tshark.sh [interface] [channel] [duration_seconds]

set -e

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

# 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
timeout "$DURATION" tshark -q -i "$INTERFACE" -n -w "$TEMP_PCAP" 2>/dev/null
CAPTURE_EXIT_CODE=$?
set -e  # Re-enable exit on error

# 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

# Now parse the pcap file to extract fields
# Use -Y filter to only process frames that have RA/TA to avoid field errors
CAPTURE_OUTPUT=$(tshark -q -r "$TEMP_PCAP" -n -T fields \
    -Y "(wlan.ra) || (wlan.ta)" \
    -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 radiotap.present \
    2>/dev/null || true)

# Clean up temp file
rm -f "$TEMP_PCAP"

# 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
PACKET_LINES=$(echo "$CAPTURE_OUTPUT" | grep -E '^[0-9]+\t' | grep -v -E '(packets captured|Capturing on|Running as)' || true)
FINAL_COUNT=$(echo "$PACKET_LINES" | wc -l || echo "0")

# Count packets with PLCP headers (radiotap present)
# radiotap.present field is the 8th field (after frame.number, frame.time, wlan.ra, wlan.ta, wlan.fc.type, wlan.fc.subtype, wlan.fc.type_subtype)
PLCP_COUNT=$(echo "$PACKET_LINES" | awk -F'\t' 'NF >= 8 && $1 != "" && $8 != "" && $8 != "0" && $8 != "-" {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"
        radiotap = ($8 == "1" || $8 == "1.0") ? "yes" : (($8 != "" && $8 != "-") ? "no" : "N/A")
        printf "  Frame %s: RA=%s, TA=%s, type=%s/%s, PLCP=%s\n",
            $1, ra, ta, type, subtype, radiotap
    }'
    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 ""
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