#!/bin/bash

# NxLite Vulnerability Testing Script
# Tests various security vulnerabilities and attack vectors

set -e

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

SERVER_URL="http://localhost:7877"
TEST_DIR="vuln_test_files"
RESULTS_FILE="vulnerability_test_results.txt"

TOTAL_TESTS=0
PASSED_TESTS=0
FAILED_TESTS=0

print_header() {
    echo -e "${BLUE}================================${NC}"
    echo -e "${BLUE}  NxLite Vulnerability Testing  ${NC}"
    echo -e "${BLUE}================================${NC}"
    echo ""
}

print_test() {
    echo -e "${YELLOW}[TEST]${NC} $1"
    TOTAL_TESTS=$((TOTAL_TESTS + 1))
}

print_pass() {
    echo -e "${GREEN}[PASS]${NC} $1"
    PASSED_TESTS=$((PASSED_TESTS + 1))
}

print_fail() {
    echo -e "${RED}[FAIL]${NC} $1"
    FAILED_TESTS=$((FAILED_TESTS + 1))
}

print_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

check_server() {
    print_info "Checking if server is running..."
    if ! curl -s --max-time 5 "$SERVER_URL" > /dev/null 2>&1; then
        echo -e "${RED}Error: Server is not running at $SERVER_URL${NC}"
        echo "Please start the server first: ./NxLite"
        exit 1
    fi
    print_info "Server is running ✓"
    echo ""
}

setup_test_env() {
    print_info "Setting up test environment..."
    mkdir -p "$TEST_DIR"
    
    echo "<!DOCTYPE html><html><body><h1>Test Page</h1></body></html>" > "$TEST_DIR/test.html"
    echo "console.log('test');" > "$TEST_DIR/test.js"
    echo "body { color: red; }" > "$TEST_DIR/test.css"
    echo "Test content" > "$TEST_DIR/test.txt"
    
    dd if=/dev/zero of="$TEST_DIR/large_file.bin" bs=1M count=10 2>/dev/null
    
    echo "" > "$RESULTS_FILE"
    print_info "Test environment ready ✓"
    echo ""
}

test_path_traversal() {
    print_test "Path Traversal Attack"
    
    local payloads=(
        "../../../etc/passwd"
        "..%2F..%2F..%2Fetc%2Fpasswd"
        "....//....//....//etc/passwd"
        "%2e%2e%2f%2e%2e%2f%2e%2e%2fetc%2fpasswd"
        "..\\..\\..\\windows\\system32\\drivers\\etc\\hosts"
    )
    
    local blocked=0
    for payload in "${payloads[@]}"; do
        response=$(curl -s --max-time 5 -o /dev/null -w "%{http_code}" "$SERVER_URL/$payload" 2>/dev/null || echo "000")
        if [[ "$response" == "403" || "$response" == "404" || "$response" == "400" ]]; then
            blocked=$((blocked + 1))
        fi
        sleep 0.1
    done
    
    if [[ $blocked -eq ${#payloads[@]} ]]; then
        print_pass "Path traversal attacks blocked"
        echo "Path Traversal: PASS - All payloads blocked" >> "$RESULTS_FILE"
    else
        print_fail "Some path traversal attacks succeeded ($blocked/${#payloads[@]} blocked)"
        echo "Path Traversal: FAIL - $blocked/${#payloads[@]} blocked" >> "$RESULTS_FILE"
    fi
}

test_header_injection() {
    print_test "HTTP Header Injection"
    
    local payloads=(
        $'\r\nX-Injected-Header: malicious'
        $'test\r\nX-Injected: value'
        $'%0d%0aX-Injected: header'
        $'%0aX-Injected: newline'
    )
    
    local safe=0
    for payload in "${payloads[@]}"; do
        response=$(curl -s --max-time 5 -H "X-Test: $payload" -I "$SERVER_URL/" 2>/dev/null | grep -i "x-injected" || true)
        if [[ -z "$response" ]]; then
            safe=$((safe + 1))
        fi
        sleep 0.1
    done
    
    if [[ $safe -eq ${#payloads[@]} ]]; then
        print_pass "Header injection attacks prevented"
        echo "Header Injection: PASS - All injections prevented" >> "$RESULTS_FILE"
    else
        print_fail "Some header injections succeeded ($safe/${#payloads[@]} safe)"
        echo "Header Injection: FAIL - $safe/${#payloads[@]} safe" >> "$RESULTS_FILE"
    fi
}

test_security_headers() {
    print_test "Security Headers Presence"
    
    response=""
    for attempt in 1 2 3; do
        response=$(curl -s --max-time 10 -D- -o /dev/null "$SERVER_URL/" 2>/dev/null || echo "")
        if [[ -n "$response" ]]; then
            break
        fi
        sleep 1
    done
    
    if [[ -z "$response" ]]; then
        response=$(timeout 10 curl -D- -o /dev/null "$SERVER_URL/" 2>/dev/null || echo "")
    fi
    
    local headers=(
        "X-Content-Type-Options"
        "X-Frame-Options"
        "X-XSS-Protection"
        "Referrer-Policy"
        "Content-Security-Policy"
    )
    
    local present=0
    for header in "${headers[@]}"; do
        if echo "$response" | grep -qi "$header"; then
            present=$((present + 1))
        fi
    done
    
    if [[ $present -eq ${#headers[@]} ]]; then
        print_pass "All security headers present"
        echo "Security Headers: PASS - All $present headers present" >> "$RESULTS_FILE"
    else
        print_fail "Missing security headers ($present/${#headers[@]} present)"
        echo "Security Headers: FAIL - $present/${#headers[@]} present" >> "$RESULTS_FILE"
        echo "Debug - Response received: '$response'" >> "$RESULTS_FILE"
    fi
}

test_large_request() {
    print_test "Large Request Handling"
    
    large_url="http://localhost:7877/$(python3 -c "print('A' * 8000)" 2>/dev/null || printf 'A%.0s' {1..8000})"
    response1=$(curl -s --max-time 10 -o /dev/null -w "%{http_code}" "$large_url" 2>/dev/null || echo "000")
    response1=$(echo "$response1" | grep -o "[0-9][0-9][0-9]" | head -1 || echo "000")
    
    large_header=$(python3 -c "print('A' * 9000)" 2>/dev/null || printf 'A%.0s' {1..9000})
    response2=$(curl -s --max-time 5 -o /dev/null -w "%{http_code}" -H "X-Large: $large_header" "$SERVER_URL/" 2>/dev/null || echo "000")
    response2=$(echo "$response2" | grep -o "[0-9][0-9][0-9]" | head -1 || echo "000")
    
    large_data=$(python3 -c "print('A' * 100000)" 2>/dev/null || echo "$(printf 'A%.0s' {1..100000})")
    response3=$(curl -s --max-time 10 -o /dev/null -w "%{http_code}" -X POST -d "$large_data" "$SERVER_URL/" 2>/dev/null || echo "000")
    response3=$(echo "$response3" | grep -o "[0-9][0-9][0-9]" | head -1 || echo "000")
    
    local rejected=0
    if [[ "$response1" == "413" || "$response1" == "400" || "$response1" == "414" || "$response1" == "505" || "$response1" == "000" ]]; then
        rejected=$((rejected + 1))
    fi
    if [[ "$response2" == "413" || "$response2" == "400" || "$response2" == "414" || "$response2" == "000" ]]; then
        rejected=$((rejected + 1))
    fi
    if [[ "$response3" == "413" || "$response3" == "400" || "$response3" == "414" || "$response3" == "501" || "$response3" == "000" ]]; then
        rejected=$((rejected + 1))
    fi
    
    if [[ $rejected -eq 3 ]]; then
        print_pass "Large requests properly rejected"
        echo "Large Request: PASS - All requests rejected (URL:$response1, Header:$response2, POST:$response3)" >> "$RESULTS_FILE"
    elif [[ $rejected -eq 2 ]]; then
        print_pass "Large requests mostly rejected"
        echo "Large Request: PASS - Most requests rejected (URL:$response1, Header:$response2, POST:$response3)" >> "$RESULTS_FILE"
    else
        print_fail "Some large requests accepted (URL:$response1, Header:$response2, POST:$response3)"
        echo "Large Request: FAIL - Some requests accepted (URL:$response1, Header:$response2, POST:$response3)" >> "$RESULTS_FILE"
    fi
}

test_malformed_requests() {
    print_test "Malformed Request Handling"
    
    local malformed_requests=(
        "INVALID_METHOD / HTTP/1.1\r\nHost: localhost\r\n\r\n"
        "GET / HTTP/999.999\r\nHost: localhost\r\n\r\n"
        "GET /\r\nHost: localhost\r\n\r\n"
        "GET / HTTP/1.1\r\nInvalid-Header-Line\r\n\r\n"
    )
    
    local handled=0
    for request in "${malformed_requests[@]}"; do
        {
            exec 3<>/dev/tcp/localhost/7877 2>/dev/null
            if [[ $? -eq 0 ]]; then
                echo -e "$request" >&3
                response=$(timeout 2 head -1 <&3 | grep -o "[0-9][0-9][0-9]" || echo "000")
                exec 3<&-
                exec 3>&-
            else
                response="000"
            fi
        } 2>/dev/null
        
        if [[ "$response" == "400" || "$response" == "000" || "$response" == "501" || "$response" == "505" ]]; then
            handled=$((handled + 1))
        fi
        sleep 0.1
    done
    
    if [[ $handled -ge 3 ]]; then
        print_pass "Malformed requests properly handled"
        echo "Malformed Requests: PASS - $handled/${#malformed_requests[@]} requests handled" >> "$RESULTS_FILE"
    else
        print_fail "Some malformed requests not handled ($handled/${#malformed_requests[@]})"
        echo "Malformed Requests: FAIL - $handled/${#malformed_requests[@]} handled" >> "$RESULTS_FILE"
    fi
}

test_information_disclosure() {
    print_test "Information Disclosure"
    
    response=$(curl -s --max-time 5 -I "$SERVER_URL/" 2>/dev/null || echo "")
    
    server_header=$(echo "$response" | grep -i "server:" || true)
    
    if [[ -z "$server_header" ]] || [[ "$server_header" != *"NxLite"* ]]; then
        print_pass "Server version not disclosed"
        echo "Information Disclosure: PASS - Server version hidden" >> "$RESULTS_FILE"
    else
        print_fail "Server version disclosed: $server_header"
        echo "Information Disclosure: FAIL - Server version exposed" >> "$RESULTS_FILE"
    fi
}

test_http_methods() {
    print_test "HTTP Method Security"
    
    local methods=("TRACE" "OPTIONS" "DELETE" "PUT" "PATCH")
    local blocked=0
    
    for method in "${methods[@]}"; do
        response=$(curl -s --max-time 5 -o /dev/null -w "%{http_code}" -X "$method" "$SERVER_URL/" 2>/dev/null || echo "000")
        if [[ "$response" == "405" || "$response" == "501" ]]; then
            blocked=$((blocked + 1))
        fi
        sleep 0.1
    done
    
    if [[ $blocked -eq ${#methods[@]} ]]; then
        print_pass "Dangerous HTTP methods blocked"
        echo "HTTP Methods: PASS - All dangerous methods blocked" >> "$RESULTS_FILE"
    else
        print_fail "Some dangerous methods allowed ($blocked/${#methods[@]} blocked)"
        echo "HTTP Methods: FAIL - $blocked/${#methods[@]} blocked" >> "$RESULTS_FILE"
    fi
}

test_file_extensions() {
    print_test "File Extension Security"
    
    local dangerous_files=(
        ".htaccess"
        ".env"
        "config.ini"
        "backup.sql"
        "database.db"
        ".git/config"
    )
    
    local blocked=0
    for file in "${dangerous_files[@]}"; do
        response=$(curl -s --max-time 5 -o /dev/null -w "%{http_code}" "$SERVER_URL/$file" 2>/dev/null || echo "000")
        if [[ "$response" == "403" || "$response" == "404" ]]; then
            blocked=$((blocked + 1))
        fi
        sleep 0.1
    done
    
    if [[ $blocked -eq ${#dangerous_files[@]} ]]; then
        print_pass "Dangerous file access blocked"
        echo "File Extensions: PASS - All dangerous files blocked" >> "$RESULTS_FILE"
    else
        print_fail "Some dangerous files accessible ($blocked/${#dangerous_files[@]} blocked)"
        echo "File Extensions: FAIL - $blocked/${#dangerous_files[@]} blocked" >> "$RESULTS_FILE"
    fi
}

test_null_byte_injection() {
    print_test "Null Byte Injection"
    
    local payloads=(
        "test.txt%00.php"
        "test.txt\x00.exe"
        "test%00.txt"
    )
    
    local safe=0
    for payload in "${payloads[@]}"; do
        response=$(curl -s --max-time 5 -o /dev/null -w "%{http_code}" "$SERVER_URL/$payload" 2>/dev/null || echo "000")
        if [[ "$response" == "400" || "$response" == "403" || "$response" == "404" ]]; then
            safe=$((safe + 1))
        fi
        sleep 0.1
    done
    
    if [[ $safe -eq ${#payloads[@]} ]]; then
        print_pass "Null byte injection prevented"
        echo "Null Byte Injection: PASS - All injections prevented" >> "$RESULTS_FILE"
    else
        print_fail "Some null byte injections succeeded ($safe/${#payloads[@]} safe)"
        echo "Null Byte Injection: FAIL - $safe/${#payloads[@]} safe" >> "$RESULTS_FILE"
    fi
}

test_response_splitting() {
    print_test "HTTP Response Splitting"
    
    local payloads=(
        $'test\r\nContent-Length: 0\r\n\r\nHTTP/1.1 200 OK\r\n'
        $'%0d%0aContent-Length:%200%0d%0a%0d%0a'
        $'test%0aSet-Cookie:%20malicious=true'
    )
    
    local safe=0
    for payload in "${payloads[@]}"; do
        response=$(curl -s --max-time 5 -I "$SERVER_URL/?param=$payload" 2>/dev/null | wc -l || echo "0")
        if [[ $response -lt 20 ]]; then
            safe=$((safe + 1))
        fi
        sleep 0.1
    done
    
    if [[ $safe -eq ${#payloads[@]} ]]; then
        print_pass "Response splitting prevented"
        echo "Response Splitting: PASS - All attempts prevented" >> "$RESULTS_FILE"
    else
        print_fail "Some response splitting succeeded ($safe/${#payloads[@]} safe)"
        echo "Response Splitting: FAIL - $safe/${#payloads[@]} safe" >> "$RESULTS_FILE"
    fi
}

cleanup() {
    print_info "Cleaning up test environment..."
    rm -rf "$TEST_DIR"
    print_info "Cleanup complete ✓"
}

print_summary() {
    echo ""
    echo -e "${BLUE}================================${NC}"
    echo -e "${BLUE}        Test Summary            ${NC}"
    echo -e "${BLUE}================================${NC}"
    echo -e "Total Tests: $TOTAL_TESTS"
    echo -e "${GREEN}Passed: $PASSED_TESTS${NC}"
    echo -e "${RED}Failed: $FAILED_TESTS${NC}"
    
    if [[ $FAILED_TESTS -eq 0 ]]; then
        echo -e "${GREEN}All vulnerability tests passed! ✓${NC}"
    else
        echo -e "${YELLOW}Some tests failed. Check results for details.${NC}"
    fi
    
    echo ""
    echo "Detailed results saved to: $RESULTS_FILE"
    echo ""
}

main() {
    print_header
    check_server
    setup_test_env
    
    echo "Starting vulnerability tests..."
    echo ""
    
    test_path_traversal
    test_header_injection
    test_security_headers
    test_large_request
    test_malformed_requests
    test_information_disclosure
    test_http_methods
    test_file_extensions
    test_null_byte_injection
    test_response_splitting
    
    cleanup
    print_summary
}

if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi 