#!/bin/bash
# ==============================================================================
# Sechpoint Wallarm Smart Deployer - Container Edition (PoC Optimized)
# ==============================================================================

# --- Styling ---
YELLOW='\033[1;33m'
GREEN='\033[0;32m'
RED='\033[0;31m'
NC='\033[0m'

LOG_FILE="/var/log/wallarm-deploy.log"
EU_DATA_NODES=("api.wallarm.com" "node-data0.eu1.wallarm.com" "node-data1.eu1.wallarm.com")
US_DATA_NODES=("us1.api.wallarm.com" "node-data0.us1.wallarm.com" "node-data1.us1.wallarm.com")

# --- Initialization ---
sudo touch "$LOG_FILE" && sudo chmod 644 "$LOG_FILE"
exec > >(tee -a "$LOG_FILE") 2>&1 # Log everything to file while showing on screen

clear
echo -e "${YELLOW}====================================================${NC}"
echo -e "${YELLOW}      Wallarm Automated Container Deployer          ${NC}"
echo -e "${YELLOW}====================================================${NC}"

# --- 1. PRE-FLIGHT FUNCTIONS ---

check_sudo() {
    echo -e "\n${YELLOW}[1/4] Checking Sudo...${NC}"
    if sudo -v; then
        echo -e "${GREEN}[PASS]${NC} Sudo access confirmed."
        return 0
    else
        echo -e "${RED}[FAIL]${NC} Sudo access denied."; return 1
    fi
}

check_wallarm_cloud() {
    echo -e "\n${YELLOW}[2/4] Testing Wallarm Cloud Connectivity (Port 443)...${NC}"
    local fail=0
    
    # We ask for cloud preference early to avoid testing everything unnecessarily
    read -p "Wallarm Cloud (US/EU) [US]: " CLOUD_SEL
    CLOUD_SEL=${CLOUD_SEL^^}
    CLOUD_SEL=${CLOUD_SEL:-US}
    
    local nodes_to_test=("${US_DATA_NODES[@]}")
    if [[ "$CLOUD_SEL" == "EU" ]]; then
        nodes_to_test=("${EU_DATA_NODES[@]}")
    fi

    echo "Testing $CLOUD_SEL Cloud Endpoints..."
    for node in "${nodes_to_test[@]}"; do
        if ! curl -skI --connect-timeout 5 "https://$node" > /dev/null 2>&1; then
            echo -e "${RED}[FAIL]${NC} Cannot reach $node"; fail=1
        else
            echo -e "${GREEN}[PASS]${NC} Reached $node"
        fi
    done
    
    API_HOST=$([[ "$CLOUD_SEL" == "EU" ]] && echo "api.wallarm.com" || echo "us1.api.wallarm.com")
    return $fail
}

# --- 2. INPUT & CONFIGURATION ---

get_user_input() {
    echo -e "\n${YELLOW}[3/4] Configuration & Workspace Setup...${NC}"
    
    # Instance ID Logic - Simplified to numeric directory structure
    echo -e "Existing Deployments in /opt/wallarm/:"
    if [ -d /opt/wallarm ]; then
        ls -F /opt/wallarm/ | grep '/' | sed 's/\///' || echo "None"
    else
        echo "None"
    fi
    echo ""
    
    read -p "Enter Instance Number (e.g., 1, 2, 3): " INSTANCE_NUM
    if ! [[ "$INSTANCE_NUM" =~ ^[0-9]+$ ]]; then
        echo -e "${RED}ERROR: Please enter a valid number.${NC}"; exit 1
    fi

    NODE_NAME="wallarm-node-$INSTANCE_NUM"
    INSTANCE_DIR="/opt/wallarm/$INSTANCE_NUM"
    TRAFFIC_PORT=$((8000 + INSTANCE_NUM))
    MONITOR_PORT=$((9000 + INSTANCE_NUM))

    # App Server Logic
    read -p "Enter Upstream IP (App Server) [127.0.0.1]: " UPSTREAM_IP
    UPSTREAM_IP=${UPSTREAM_IP:-127.0.0.1}
    read -p "Enter Upstream Port [80]: " UPSTREAM_PORT
    UPSTREAM_PORT=${UPSTREAM_PORT:-80}

    read -p "Paste Wallarm Token ($CLOUD_SEL Cloud): " TOKEN
    
    echo -n "Verifying connection to App Server ($UPSTREAM_IP:$UPSTREAM_PORT)... "
    if ! timeout 2 bash -c "cat < /dev/null > /dev/tcp/$UPSTREAM_IP/$UPSTREAM_PORT" 2>/dev/null; then
        echo -e "${RED}FAILED${NC}"
        echo -e "${RED}❌ ERROR: VM cannot reach internal app server at $UPSTREAM_IP:$UPSTREAM_PORT.${NC}"; exit 1
    else
        echo -e "${GREEN}OK${NC}"
    fi
}

# --- 3. ENGINE SETUP ---

setup_engine() {
    echo -e "\n${YELLOW}[4/4] 🛠️ Ensuring Engine (Podman/Docker) is ready...${NC}"
    if [ -f /etc/redhat-release ]; then
        ENGINE="podman"
        echo "Detected RHEL/CentOS. Setting up Podman..."
        sudo dnf install -y epel-release podman podman-docker wget curl &>/dev/null
        sudo systemctl enable --now podman.socket &>/dev/null
        sudo firewall-cmd --permanent --add-port=$TRAFFIC_PORT/tcp --add-port=$MONITOR_PORT/tcp &>/dev/null
        sudo firewall-cmd --reload &>/dev/null
    else
        ENGINE="docker"
        echo "Detected Ubuntu/Debian. Setting up Docker..."
        sudo apt update && sudo apt install -y docker.io wget curl &>/dev/null
        sudo systemctl enable --now docker &>/dev/null
    fi
    
    if ! command -v docker-compose &> /dev/null && ! command -v podman-compose &> /dev/null; then
        echo "Installing Compose utility..."
        if [ "$ENGINE" == "docker" ]; then sudo apt install -y docker-compose &>/dev/null; fi
        if [ "$ENGINE" == "podman" ]; then sudo dnf install -y podman-compose &>/dev/null; fi
    fi
}

# --- 4. DEPLOYMENT ---

execute_deployment() {
    echo -e "\n${YELLOW}🚀 Preparing Workspace: $INSTANCE_DIR${NC}"
    sudo mkdir -p "$INSTANCE_DIR"
    cd "$INSTANCE_DIR"

    # Fully qualified name ensures Podman/Docker doesn't prompt for registry choice
    IMAGE_NAME="docker.io/wallarm/node:latest"

    echo "Generating Nginx Configuration..."
    sudo tee "$INSTANCE_DIR/nginx.conf" > /dev/null <<EOF
server {
    listen 80;
    wallarm_mode monitoring; # PoC Safety Mode
    
    location / {
        proxy_pass http://$UPSTREAM_IP:$UPSTREAM_PORT;
        proxy_set_header Host \$host;
        proxy_set_header X-Real-IP \$remote_addr;
        proxy_set_header X-Forwarded-For \$proxy_add_x_forwarded_for;
    }
}
server { 
    listen 90; 
    location /wallarm-status { 
        wallarm_status on; 
        allow all; 
    } 
}
EOF

    echo "Generating Deployment Manifest (compose.yml)..."
    sudo tee "$INSTANCE_DIR/compose.yml" > /dev/null <<EOF
version: '3.8'
services:
  node:
    image: $IMAGE_NAME
    container_name: $NODE_NAME
    restart: always
    ports: 
      - "$TRAFFIC_PORT:80"
      - "$MONITOR_PORT:90"
    environment:
      - WALLARM_API_TOKEN=$TOKEN
      - WALLARM_API_HOST=$API_HOST
    volumes: 
      - ./nginx.conf:/etc/nginx/http.d/default.conf:ro,Z
EOF

    echo -e "${YELLOW}🚀 Launching Instance $INSTANCE_NUM ($NODE_NAME)...${NC}"
    sudo $ENGINE rm -f "$NODE_NAME" &>/dev/null

    # Pulling explicitly with docker.io prefix to avoid short-name resolution errors
    echo "Pulling latest image from Docker Hub (docker.io)..."
    sudo $ENGINE pull $IMAGE_NAME

    if command -v podman-compose &> /dev/null; then
        sudo podman-compose -f compose.yml up -d
    else
        sudo docker-compose -f compose.yml up -d
    fi
}

# --- 5. VERIFICATION & ATTACK TEST ---

verify_health() {
    echo -e "\n${YELLOW}⏳ Waiting 20s for handshake and sync...${NC}"
    sleep 20
    
    echo -en "Checking instance status page (port $MONITOR_PORT)... "
    if curl -s "http://localhost:$MONITOR_PORT/wallarm-status" | grep -q "requests"; then
        echo -e "${GREEN}SUCCESS${NC}"
    else
        echo -e "${RED}WARNING: Status page not responding yet.${NC}"
        echo -e "Check logs with: sudo $ENGINE logs $NODE_NAME"
    fi

    echo -e "\n${YELLOW}⚔️ Running Attack Test (SQLi & XSS)...${NC}"
    
    # Test 1: SQL Injection
    echo -n "Sending SQLi payload to port $TRAFFIC_PORT... "
    local sqli_res=$(curl -s -o /dev/null -w "%{http_code}" "http://localhost:$TRAFFIC_PORT/?id='OR+1=1+UNION+SELECT+1,2,3--")
    echo -e "HTTP Status: $sqli_res (Logged)"

    # Test 2: XSS
    echo -n "Sending XSS payload to port $TRAFFIC_PORT... "
    local xss_res=$(curl -s -o /dev/null -w "%{http_code}" "http://localhost:$TRAFFIC_PORT/?search=<script>alert('Wallarm_Test')</script>")
    echo -e "HTTP Status: $xss_res (Logged)"

    echo -e "\n${GREEN}✅ DEPLOYMENT FINISHED${NC}"
    echo -e "--------------------------------------------------"
    echo -e "Instance ID:    $INSTANCE_NUM"
    echo -e "Traffic Port:   $TRAFFIC_PORT"
    echo -e "Monitor Port:   $MONITOR_PORT"
    echo -e "\nCheck your Wallarm Console ($CLOUD_SEL Cloud) now."
    echo -e "The attacks should appear in the 'Events' section within 1-2 minutes."
    echo -e "--------------------------------------------------"
}

# --- MAIN FLOW ---
check_sudo || exit 1
check_wallarm_cloud || { echo -e "${RED}Cloud connectivity failed. Cannot continue.${NC}"; exit 1; }
get_user_input
setup_engine
execute_deployment
verify_health