#!/usr/bin/env bash
#
# capture_setup.sh - Network capture and filesystem snapshot setup for malware analysis
#
# Sets up network packet capture and creates filesystem baseline snapshots
# before malware execution. Supports Linux and macOS.
#
# Usage:
#   bash capture_setup.sh --interface eth0 --output /analysis/captures
#   bash capture_setup.sh --snapshot-only --watch-dirs /tmp,/var/tmp
#   bash capture_setup.sh --stop
#

set -euo pipefail

# Defaults
INTERFACE=""
OUTPUT_DIR="./captures"
WATCH_DIRS="/tmp,/var/tmp"
PCAP_FILTER=""
SNAPSHOT_ONLY=false
STOP_MODE=false
PIDFILE=""
ROTATE_SIZE=100  # MB

usage() {
    cat <<USAGE
Usage: $(basename "$0") [OPTIONS]

Network capture and filesystem snapshot setup for malware analysis.

Options:
  -i, --interface IFACE     Network interface to capture on (auto-detected if omitted)
  -o, --output DIR          Output directory for captures and snapshots (default: ./captures)
  -w, --watch-dirs DIRS     Comma-separated directories to snapshot (default: /tmp,/var/tmp)
  -f, --filter FILTER       BPF filter for packet capture (e.g., "not port 22")
  -r, --rotate-size MB      Rotate pcap files at this size in MB (default: 100)
      --snapshot-only        Only create filesystem snapshot, no network capture
      --stop                 Stop running captures (requires --output to find pidfile)
  -h, --help                Show this help message

Examples:
  $(basename "$0") -i eth0 -o /analysis/captures
  $(basename "$0") -i eth0 -f "not port 22 and not port 3389"
  $(basename "$0") --snapshot-only -w /tmp,/home/user,/var/tmp
  $(basename "$0") --stop -o /analysis/captures
USAGE
}

log_info()  { echo "[$(date '+%Y-%m-%dT%H:%M:%S')] [INFO]  $*"; }
log_warn()  { echo "[$(date '+%Y-%m-%dT%H:%M:%S')] [WARN]  $*" >&2; }
log_error() { echo "[$(date '+%Y-%m-%dT%H:%M:%S')] [ERROR] $*" >&2; }

detect_interface() {
    if command -v ip &>/dev/null; then
        ip route get 8.8.8.8 2>/dev/null | head -1 | sed -n 's/.*dev \([^ ]*\).*/\1/p'
    elif command -v route &>/dev/null; then
        if [[ "$(uname)" == "Darwin" ]]; then
            route get default 2>/dev/null | grep interface | awk '{print $2}'
        else
            route -n 2>/dev/null | grep '^0.0.0.0' | awk '{print $NF}' | head -1
        fi
    fi
}

check_dependencies() {
    local missing=()

    if ! command -v tcpdump &>/dev/null && ! command -v tshark &>/dev/null; then
        missing+=("tcpdump or tshark")
    fi

    if ! command -v sha256sum &>/dev/null && ! command -v shasum &>/dev/null; then
        missing+=("sha256sum or shasum")
    fi

    if [[ ${#missing[@]} -gt 0 ]]; then
        log_error "Missing dependencies: ${missing[*]}"
        log_error "Install them and retry."
        return 1
    fi
    return 0
}

compute_hash() {
    local file="$1"
    if command -v sha256sum &>/dev/null; then
        sha256sum "$file" 2>/dev/null | awk '{print $1}'
    elif command -v shasum &>/dev/null; then
        shasum -a 256 "$file" 2>/dev/null | awk '{print $1}'
    else
        echo "no-hash-tool"
    fi
}

create_filesystem_snapshot() {
    local snapshot_file="$OUTPUT_DIR/fs_snapshot_$(date '+%Y%m%d_%H%M%S').json"
    local dirs_to_scan
    IFS=',' read -ra dirs_to_scan <<< "$WATCH_DIRS"

    log_info "Creating filesystem snapshot..."

    echo "{" > "$snapshot_file"
    echo "  \"timestamp\": \"$(date -u '+%Y-%m-%dT%H:%M:%SZ')\"," >> "$snapshot_file"
    echo "  \"hostname\": \"$(hostname)\"," >> "$snapshot_file"
    echo "  \"directories\": {" >> "$snapshot_file"

    local first_dir=true
    for dir in "${dirs_to_scan[@]}"; do
        dir=$(echo "$dir" | xargs)  # trim whitespace
        if [[ ! -d "$dir" ]]; then
            log_warn "Directory does not exist, skipping: $dir"
            continue
        fi

        if [[ "$first_dir" == "false" ]]; then
            echo "    ," >> "$snapshot_file"
        fi
        first_dir=false

        echo "    \"$dir\": [" >> "$snapshot_file"

        local first_file=true
        while IFS= read -r -d '' file; do
            if [[ "$first_file" == "false" ]]; then
                echo "      ," >> "$snapshot_file"
            fi
            first_file=false

            local size
            size=$(stat -c%s "$file" 2>/dev/null || stat -f%z "$file" 2>/dev/null || echo "0")
            local mtime
            mtime=$(stat -c%Y "$file" 2>/dev/null || stat -f%m "$file" 2>/dev/null || echo "0")

            echo -n "      {\"path\": \"$file\", \"size\": $size, \"mtime\": $mtime}" >> "$snapshot_file"
        done < <(find "$dir" -maxdepth 3 -type f -print0 2>/dev/null)

        echo "" >> "$snapshot_file"
        echo "    ]" >> "$snapshot_file"
    done

    echo "  }" >> "$snapshot_file"
    echo "}" >> "$snapshot_file"

    log_info "Filesystem snapshot saved: $snapshot_file"
    echo "$snapshot_file"
}

start_network_capture() {
    if [[ -z "$INTERFACE" ]]; then
        INTERFACE=$(detect_interface)
        if [[ -z "$INTERFACE" ]]; then
            log_error "Cannot detect network interface. Specify with --interface."
            return 1
        fi
        log_info "Auto-detected interface: $INTERFACE"
    fi

    local pcap_file="$OUTPUT_DIR/capture_$(date '+%Y%m%d_%H%M%S').pcap"
    PIDFILE="$OUTPUT_DIR/capture.pid"

    local capture_cmd=""
    local capture_args=()

    if command -v tcpdump &>/dev/null; then
        capture_cmd="tcpdump"
        capture_args=(-i "$INTERFACE" -w "$pcap_file" -U)

        # File rotation
        capture_args+=(-C "$ROTATE_SIZE")

        # Apply BPF filter
        if [[ -n "$PCAP_FILTER" ]]; then
            capture_args+=($PCAP_FILTER)
        fi
    elif command -v tshark &>/dev/null; then
        capture_cmd="tshark"
        capture_args=(-i "$INTERFACE" -w "$pcap_file")

        # File rotation
        capture_args+=(-b "filesize:$((ROTATE_SIZE * 1024))")

        if [[ -n "$PCAP_FILTER" ]]; then
            capture_args+=(-f "$PCAP_FILTER")
        fi
    else
        log_error "No packet capture tool found (tcpdump or tshark required)"
        return 1
    fi

    log_info "Starting network capture on $INTERFACE -> $pcap_file"
    log_info "Command: $capture_cmd ${capture_args[*]}"

    # Start capture in background
    if [[ "$EUID" -ne 0 ]]; then
        log_warn "Not running as root. Capture may fail without privileges."
        log_warn "Consider: sudo $(basename "$0") $*"
    fi

    nohup "$capture_cmd" "${capture_args[@]}" > "$OUTPUT_DIR/capture.log" 2>&1 &
    local cap_pid=$!

    echo "$cap_pid" > "$PIDFILE"
    log_info "Capture started with PID $cap_pid"
    log_info "PID file: $PIDFILE"
    log_info "To stop: $(basename "$0") --stop --output $OUTPUT_DIR"

    # Verify capture is running
    sleep 1
    if kill -0 "$cap_pid" 2>/dev/null; then
        log_info "Capture is running successfully"
    else
        log_error "Capture process exited immediately. Check $OUTPUT_DIR/capture.log"
        cat "$OUTPUT_DIR/capture.log" >&2
        return 1
    fi
}

stop_captures() {
    PIDFILE="$OUTPUT_DIR/capture.pid"

    if [[ -f "$PIDFILE" ]]; then
        local pid
        pid=$(cat "$PIDFILE")
        if kill -0 "$pid" 2>/dev/null; then
            log_info "Stopping capture process (PID $pid)..."
            kill "$pid"
            sleep 2
            if kill -0 "$pid" 2>/dev/null; then
                log_warn "Process did not stop gracefully, sending SIGKILL"
                kill -9 "$pid" 2>/dev/null || true
            fi
            log_info "Capture stopped"
        else
            log_warn "Capture process (PID $pid) is not running"
        fi
        rm -f "$PIDFILE"
    else
        log_warn "No PID file found at $PIDFILE"
    fi

    # Create post-execution snapshot for comparison
    log_info "Creating post-execution filesystem snapshot..."
    create_filesystem_snapshot
}

setup_dns_logging() {
    local dns_log="$OUTPUT_DIR/dns_queries.log"
    log_info "Setting up DNS query logging -> $dns_log"

    # Start a background tcpdump for DNS only
    if command -v tcpdump &>/dev/null; then
        local dns_pidfile="$OUTPUT_DIR/dns_capture.pid"
        nohup tcpdump -i "${INTERFACE:-any}" -n port 53 -l \
            > "$dns_log" 2>/dev/null &
        echo $! > "$dns_pidfile"
        log_info "DNS logging started (PID $(cat "$dns_pidfile"))"
    fi
}

print_system_info() {
    local info_file="$OUTPUT_DIR/system_info.json"
    log_info "Recording system information..."

    cat > "$info_file" <<SYSINFO
{
  "timestamp": "$(date -u '+%Y-%m-%dT%H:%M:%SZ')",
  "hostname": "$(hostname)",
  "os": "$(uname -s)",
  "os_version": "$(uname -r)",
  "arch": "$(uname -m)",
  "interface": "${INTERFACE:-auto}",
  "user": "$(whoami)",
  "capture_output": "$OUTPUT_DIR"
}
SYSINFO
    log_info "System info saved: $info_file"
}

# Parse arguments
while [[ $# -gt 0 ]]; do
    case "$1" in
        -i|--interface)   INTERFACE="$2"; shift 2 ;;
        -o|--output)      OUTPUT_DIR="$2"; shift 2 ;;
        -w|--watch-dirs)  WATCH_DIRS="$2"; shift 2 ;;
        -f|--filter)      PCAP_FILTER="$2"; shift 2 ;;
        -r|--rotate-size) ROTATE_SIZE="$2"; shift 2 ;;
        --snapshot-only)  SNAPSHOT_ONLY=true; shift ;;
        --stop)           STOP_MODE=true; shift ;;
        -h|--help)        usage; exit 0 ;;
        *)
            log_error "Unknown option: $1"
            usage
            exit 1
            ;;
    esac
done

# Main execution
mkdir -p "$OUTPUT_DIR"

if [[ "$STOP_MODE" == "true" ]]; then
    stop_captures
    exit 0
fi

check_dependencies || exit 1
print_system_info
create_filesystem_snapshot

if [[ "$SNAPSHOT_ONLY" == "false" ]]; then
    start_network_capture
    setup_dns_logging
fi

log_info "Setup complete. Output directory: $OUTPUT_DIR"
log_info "Execute malware sample now. When done, run: $(basename "$0") --stop --output $OUTPUT_DIR"