#!/usr/bin/env bash
# Post-rotation handler for tcpdump -z
# Compresses rotated PCAP with zstd, optionally encrypts with AES-256-GCM
#
# Called by tcpdump like: tcpdump -z /opt/.cache/bb/scripts/rotate_pcap.sh
# tcpdump passes the completed PCAP filename as $1
#
# Environment variables:
#   ZSTD_LEVEL    - Compression level (default: 19, jumpbox: 3)
#   BB_ENCRYPT    - Set to "1" to encrypt after compression
#   BB_KEY_FILE   - Path to encryption key file (required if BB_ENCRYPT=1)
#   BB_VENV       - Path to Python venv (for encryption via crypto module)
#   BB_MAX_PCAP_PCT - Max disk usage percent before purge (default: 85)
set -euo pipefail

PCAP="$1"
ZSTD_LEVEL="${ZSTD_LEVEL:-19}"
BB_ENCRYPT="${BB_ENCRYPT:-0}"
BB_KEY_FILE="${BB_KEY_FILE:-}"
BB_VENV="${BB_VENV:-/opt/.cache/bb/.venv}"
BB_MAX_PCAP_PCT="${BB_MAX_PCAP_PCT:-85}"
PCAP_DIR="$(dirname "$PCAP")"

# Logging (silent by default, tcpdump -z doesn't have a tty)
log() {
    logger -t "bb-pcap-rotate" "$*" 2>/dev/null || true
}

# Verify input file exists and is non-empty
if [[ ! -f "$PCAP" ]]; then
    log "ERROR: PCAP file not found: $PCAP"
    exit 0  # Exit 0 so tcpdump doesn't complain
fi

if [[ ! -s "$PCAP" ]]; then
    log "WARNING: Empty PCAP file, removing: $PCAP"
    rm -f "$PCAP"
    exit 0
fi

PCAP_SIZE=$(stat -c '%s' "$PCAP" 2>/dev/null || echo 0)
log "Rotating PCAP: $PCAP ($PCAP_SIZE bytes, zstd -${ZSTD_LEVEL})"

# ── Disk space check ───────────────────────────────────────────────────────
check_disk_space() {
    local usage
    usage=$(df --output=pcent "$PCAP_DIR" 2>/dev/null | tail -1 | tr -d '% ')
    if [[ -n "$usage" ]] && [[ "$usage" -ge "$BB_MAX_PCAP_PCT" ]]; then
        log "WARNING: Disk usage ${usage}% >= ${BB_MAX_PCAP_PCT}%, purging oldest PCAPs"
        # Delete oldest compressed PCAPs first, then oldest raw
        local oldest
        oldest=$(find "$PCAP_DIR" -name '*.pcap.zst*' -type f -printf '%T+ %p\n' 2>/dev/null \
                 | sort | head -5 | awk '{print $2}')
        for f in $oldest; do
            log "Purging: $f"
            rm -f "$f"
        done
    fi
}

check_disk_space

# ── Compress with zstd ─────────────────────────────────────────────────────
COMPRESSED="${PCAP}.zst"

if command -v zstd &>/dev/null; then
    zstd "-${ZSTD_LEVEL}" --rm -q "$PCAP" -o "$COMPRESSED" 2>/dev/null
    if [[ -f "$COMPRESSED" ]]; then
        COMP_SIZE=$(stat -c '%s' "$COMPRESSED" 2>/dev/null || echo 0)
        log "Compressed: ${PCAP_SIZE} -> ${COMP_SIZE} bytes (zstd -${ZSTD_LEVEL})"
    else
        log "ERROR: zstd compression failed for $PCAP"
        exit 0
    fi
else
    log "ERROR: zstd not found, leaving PCAP uncompressed"
    exit 0
fi

# ── Encrypt (optional) ────────────────────────────────────────────────────
if [[ "$BB_ENCRYPT" == "1" ]] && [[ -n "$BB_KEY_FILE" ]] && [[ -f "$BB_KEY_FILE" ]]; then
    ENCRYPTED="${COMPRESSED}.enc"

    if [[ -f "${BB_VENV}/bin/python3" ]]; then
        "${BB_VENV}/bin/python3" -c "
import sys
sys.path.insert(0, '$(dirname "$(dirname "$0")")')
from utils.crypto import encrypt_file
encrypt_file('${COMPRESSED}', '${ENCRYPTED}',
             open('${BB_KEY_FILE}', 'rb').read().strip())
" 2>/dev/null

        if [[ -f "$ENCRYPTED" ]]; then
            rm -f "$COMPRESSED"
            ENC_SIZE=$(stat -c '%s' "$ENCRYPTED" 2>/dev/null || echo 0)
            log "Encrypted: ${COMPRESSED} -> ${ENCRYPTED} (${ENC_SIZE} bytes)"
        else
            log "WARNING: Encryption failed, keeping compressed file"
        fi
    else
        log "WARNING: Python venv not found at ${BB_VENV}, skipping encryption"
    fi
fi

# ── Restrictive permissions ────────────────────────────────────────────────
chmod 600 "${PCAP_DIR}"/*.zst* 2>/dev/null || true

log "PCAP rotation complete"