#!/bin/bash

# Script to launch wstunnel server using parameters from the [server] section of a configuration file

# Configuration file paths
PRIMARY_CONFIG="/etc/wstunnel/wstunnel.conf"
FALLBACK_CONFIG="/usr/share/defaults/etc/wstunnel/wstunnel.conf"

# Determine which configuration file to use
CONFIG_FILE=""
if [ -f "$PRIMARY_CONFIG" ]; then
    CONFIG_FILE="$PRIMARY_CONFIG"
elif [ -f "$FALLBACK_CONFIG" ]; then
    CONFIG_FILE="$FALLBACK_CONFIG"
else
    echo "Error: Configuration file not found at '$PRIMARY_CONFIG' or '$FALLBACK_CONFIG'."
    exit 1
fi

# Check if a configuration file is provided as an argument (overrides default paths)
if [ $# -eq 1 ]; then
    CONFIG_FILE="$1"
    if [ ! -f "$CONFIG_FILE" ]; then
        echo "Error: Specified configuration file '$CONFIG_FILE' not found."
        exit 1
    fi
fi

# Initialize variables for wstunnel server parameters
BIND_ADDRESS=""
SOCKET_SO_MARK=""
WEBSOCKET_PING_FREQUENCY=""
NO_COLOR=""
WEBSOCKET_MASK_FRAME=""
DNS_RESOLVER=()
DNS_RESOLVER_PREFER_IPV4=""
LOG_LEVEL=""
RESTRICT_TO=()
RESTRICT_HTTP_UPGRADE_PATH_PREFIX=()
RESTRICT_CONFIG=""
TLS_CERTIFICATE=""
TLS_PRIVATE_KEY=""
TLS_CLIENT_CA_CERTS=""
HTTP_PROXY=""
HTTP_PROXY_LOGIN=""
HTTP_PROXY_PASSWORD=""
REMOTE_TO_LOCAL_SERVER_IDLE_TIMEOUT=""
NB_WORKER_THREADS=""

# Function to trim whitespace
trim() {
    local var="$1"
    var="${var#"${var%%[![:space:]]*}"}" # Remove leading whitespace
    var="${var%"${var##*[![:space:]]}"}" # Remove trailing whitespace
    echo -n "$var"
}

# Parse the [server] section of the INI file
current_section=""
while IFS='=' read -r key value; do
    # Skip empty lines and comments
    if [[ -z "$key" || "$key" =~ ^\s*# || "$key" =~ ^\s*\; ]]; then
        continue
    fi

    # Check for section headers
    if [[ "$key" =~ ^\s*\[.*\]\s*$ ]]; then
        current_section=$(echo "$key" | sed 's/^\s*\[\(.*\)\]\s*$/\1/')
        continue
    fi

    # Process only the [server] section
    if [ "$current_section" != "server" ]; then
        continue
    fi

    # Trim whitespace from key and value
    key=$(trim "$key")
    value=$(trim "$value")

    # Skip if value is empty
    if [ -z "$value" ]; then
        continue
    fi

    # Map INI keys to wstunnel server parameters
    case "$key" in
        bind_address)
            BIND_ADDRESS="$value"
            ;;
        socket_so_mark)
            SOCKET_SO_MARK="$value"
            ;;
        websocket_ping_frequency)
            WEBSOCKET_PING_FREQUENCY="$value"
            ;;
        no_color)
            NO_COLOR="$value"
            ;;
        websocket_mask_frame)
            WEBSOCKET_MASK_FRAME="$value"
            ;;
        dns_resolver)
            DNS_RESOLVER+=("$value")
            ;;
        dns_resolver_prefer_ipv4)
            DNS_RESOLVER_PREFER_IPV4="$value"
            ;;
        log_level)
            LOG_LEVEL="$value"
            ;;
        restrict_to)
            # Split comma-separated values into array
            IFS=',' read -ra restrict_array <<< "$value"
            for restrict in "${restrict_array[@]}"; do
                RESTRICT_TO+=("$(trim "$restrict")")
            done
            ;;
        restrict_http_upgrade_path_prefix)
            # Split comma-separated values into array
            IFS=',' read -ra prefix_array <<< "$value"
            for prefix in "${prefix_array[@]}"; do
                RESTRICT_HTTP_UPGRADE_PATH_PREFIX+=("$(trim "$prefix")")
            done
            ;;
        restrict_config)
            RESTRICT_CONFIG="$value"
            ;;
        tls_certificate)
            TLS_CERTIFICATE="$value"
            ;;
        tls_private_key)
            TLS_PRIVATE_KEY="$value"
            ;;
        tls_client_ca_certs)
            TLS_CLIENT_CA_CERTS="$value"
            ;;
        http_proxy)
            HTTP_PROXY="$value"
            ;;
        http_proxy_login)
            HTTP_PROXY_LOGIN="$value"
            ;;
        http_proxy_password)
            HTTP_PROXY_PASSWORD="$value"
            ;;
        remote_to_local_server_idle_timeout)
            REMOTE_TO_LOCAL_SERVER_IDLE_TIMEOUT="$value"
            ;;
        nb_worker_threads)
            NB_WORKER_THREADS="$value"
            ;;
    esac
done < "$CONFIG_FILE"

# Build the wstunnel server command
CMD=("wstunnel" "server")

# Add bind address (required argument)
if [ -z "$BIND_ADDRESS" ]; then
    echo "Error: bind_address is required in the [server] section of the configuration file."
    exit 1
fi
CMD+=("$BIND_ADDRESS")

# Add optional parameters
[ -n "$SOCKET_SO_MARK" ] && CMD+=("--socket-so-mark" "$SOCKET_SO_MARK")
[ -n "$WEBSOCKET_PING_FREQUENCY" ] && CMD+=("--websocket-ping-frequency" "$WEBSOCKET_PING_FREQUENCY")
[ "$NO_COLOR" = "true" ] && CMD+=("--no-color" "true")
[ "$WEBSOCKET_MASK_FRAME" = "true" ] && CMD+=("--websocket-mask-frame")
for resolver in "${DNS_RESOLVER[@]}"; do
    CMD+=("--dns-resolver" "$resolver")
done
[ "$DNS_RESOLVER_PREFER_IPV4" = "true" ] && CMD+=("--dns-resolver-prefer-ipv4")
[ -n "$LOG_LEVEL" ] && CMD+=("--log-lvl" "$LOG_LEVEL")
for restrict in "${RESTRICT_TO[@]}"; do
    CMD+=("--restrict-to" "$restrict")
done
for prefix in "${RESTRICT_HTTP_UPGRADE_PATH_PREFIX[@]}"; do
    CMD+=("--restrict-http-upgrade-path-prefix" "$prefix")
done
[ -n "$RESTRICT_CONFIG" ] && CMD+=("--restrict-config" "$RESTRICT_CONFIG")
[ -n "$TLS_CERTIFICATE" ] && CMD+=("--tls-certificate" "$TLS_CERTIFICATE")
[ -n "$TLS_PRIVATE_KEY" ] && CMD+=("--tls-private-key" "$TLS_PRIVATE_KEY")
[ -n "$TLS_CLIENT_CA_CERTS" ] && CMD+=("--tls-client-ca-certs" "$TLS_CLIENT_CA_CERTS")
[ -n "$HTTP_PROXY" ] && CMD+=("--http-proxy" "$HTTP_PROXY")
[ -n "$HTTP_PROXY_LOGIN" ] && CMD+=("--http-proxy-login" "$HTTP_PROXY_LOGIN")
[ -n "$HTTP_PROXY_PASSWORD" ] && CMD+=("--http-proxy-password" "$HTTP_PROXY_PASSWORD")
[ -n "$REMOTE_TO_LOCAL_SERVER_IDLE_TIMEOUT" ] && CMD+=("--remote-to-local-server-idle-timeout" "$REMOTE_TO_LOCAL_SERVER_IDLE_TIMEOUT")

# Set environment variable for nb_worker_threads if specified
if [ -n "$NB_WORKER_THREADS" ]; then
    export TOKIO_WORKER_THREADS="$NB_WORKER_THREADS"
fi

# Print the command for debugging
echo "Using configuration file: $CONFIG_FILE"

# Execute the wstunnel server command
exec "${CMD[@]}"