#!/bin/bash

set -e

# --- Configuration & Defaults ---
DEFAULT_PORT="22"
CURL_OPTS="-fsSL"
CA_PUB_PATH="/etc/ssh/sshteam_ca.pub"
SSHD_CONFIG="/etc/ssh/sshd_config"

# --- HELPER: Redirect logs to stderr ---
info() { echo -e "\033[0;32m[INFO]\033[0m $1" >&2; }
warn() { echo -e "\033[0;33m[WARN]\033[0m $1" >&2; }
error() { echo -e "\033[0;31m[ERROR]\033[0m $1" >&2; exit 1; }

show_help() {
        cat << 'EOF'
Usage: join.sh [options] <tenant>

Options:
    -p, --port <port>     SSH port to probe (default: 22)
    -n, --name <name>     Server name override (default: hostname)
    -t, --tag <tag>       Optional server tag
    -k, --ignore-ssl-cert Skip TLS verification
    -h, --help            Show this help text

Examples:
    ./join.sh tenant.com
    ./join.sh -p 2222 tenant.com
EOF
}

# --- 1. Argument Parsing ---
POSITIONAL_ARGS=()
CMD_PORT=""
CMD_NAME=""
CMD_TAG=""

while [[ "$#" -gt 0 ]]; do
    case $1 in
        -p|--port) CMD_PORT="$2"; shift 2 ;;
        -n|--name) CMD_NAME="$2"; shift 2 ;;
        -t|--tag) CMD_TAG="$2"; shift 2 ;;
        -k|--ignore-ssl-cert) CURL_OPTS="${CURL_OPTS}k"; shift ;;
        -h|--help) show_help; exit 0 ;;
        -*) show_help; error "Unknown option: $1" ;;
        *) POSITIONAL_ARGS+=("$1"); shift ;;
    esac
done

TENANT_TO_USE=""
if [ ${#POSITIONAL_ARGS[@]} -gt 0 ]; then
    TENANT_TO_USE="${POSITIONAL_ARGS[${#POSITIONAL_ARGS[@]}-1]}"
else
    TENANT_TO_USE="$TENANT"
fi

if [ -z "$TENANT_TO_USE" ]; then
    show_help
    exit 0
fi

PORT_TO_USE="${CMD_PORT:-$DEFAULT_PORT}"

# --- 2. Metadata Collection ---
info "Collecting server metadata..."
SVR_HOSTNAME=$(hostname)
SVR_NAME="${CMD_NAME:-$SVR_HOSTNAME}"

if [[ "$OSTYPE" == "darwin"* ]]; then
    IP_ADDR=$(ifconfig | grep "inet " | grep -v 127.0.0.1 | awk '{print $2}' | head -n 1)
    OS_FAMILY="macOS"
    OS_VERSION=$(sw_vers -productVersion)
else
    IP_ADDR=$(hostname -I | awk '{print $1}' || ip addr show | grep -m1 'inet ' | awk '{print $2}' | cut -d/ -f1 || echo "unknown")
    if [ -f /etc/os-release ]; then
        OS_FAMILY=$(grep -E '^ID=' /etc/os-release | cut -d= -f2 | tr -d '"')
        OS_VERSION=$(grep -E '^VERSION_ID=' /etc/os-release | cut -d= -f2 | tr -d '"')
    else
        OS_FAMILY="linux-generic"
        OS_VERSION="unknown"
    fi
fi

# --- 3. Device Code Authorization ---
info "Requesting device authorization..."
RESP=$(curl $CURL_OPTS -s -X POST "https://$TENANT_TO_USE/oauth2/device" \
    -d "client_id=sshteam" \
    -d "scope=sshteam-ca-authority")

USER_CODE=$(echo "$RESP" | jq -r '.user_code')
VERIFY_URI=$(echo "$RESP" | jq -r '.verification_uri_complete')
DEVICE_CODE=$(echo "$RESP" | jq -r '.device_code')
INTERVAL=$(echo "$RESP" | jq -r '.interval')
[ -z "$INTERVAL" ] && INTERVAL=5

echo -e "\nACTION REQUIRED: Please visit: $VERIFY_URI"
echo "Confirm code: $USER_CODE"
read -p "Press [Enter] after approving..."

info "Waiting for authorization..."
ACCESS_TOKEN=""
while true; do
    TOKEN_RESP=$(curl $CURL_OPTS -s -X POST "https://$TENANT_TO_USE/oauth2/token" \
        -d "grant_type=urn:ietf:params:oauth:grant-type:device_code" \
        -d "device_code=$DEVICE_CODE" \
        -d "client_id=sshteam")

    ERR=$(echo "$TOKEN_RESP" | jq -r '.error')
    ACCESS_TOKEN=$(echo "$TOKEN_RESP" | jq -r '.access_token')

    if [ -n "$ACCESS_TOKEN" ] && [ "$ACCESS_TOKEN" != "null" ]; then
        break
    fi

    if [ "$ERR" == "authorization_pending" ] || [ "$ERR" == "slow_down" ]; then
        sleep "$INTERVAL"
        continue
    fi

    error "Failed to retrieve access token: $TOKEN_RESP"
done

# --- 4. Fingerprint Collection ---
info "Probing local SSH host keys..."
RAW_SCAN=$(ssh-keyscan -p "$PORT_TO_USE" localhost 2>/dev/null) || true
[ -z "$RAW_SCAN" ] && error "Could not probe local SSH keys on port $PORT_TO_USE."

# --- 5. Join and Fetch CA Keys ---
JOIN_URL="https://$TENANT_TO_USE/ssh/join"
info "Registering server '$SVR_NAME' ($IP_ADDR) with $TENANT_TO_USE..."
CURL_CMD=(curl $CURL_OPTS -X POST "$JOIN_URL"
    -H "Authorization: Bearer $ACCESS_TOKEN"
    --data-urlencode "port=$PORT_TO_USE"
    --data-urlencode "name=$SVR_NAME"
    --data-urlencode "hostname=$SVR_HOSTNAME"
    --data-urlencode "ipAddress=$IP_ADDR"
    --data-urlencode "osFamily=$OS_FAMILY"
    --data-urlencode "osVersion=$OS_VERSION")

if [ -n "$CMD_TAG" ]; then
    CURL_CMD+=(--data-urlencode "tag=$CMD_TAG")
fi

FP_LIST=$(while read -r line; do
    if [ -n "$line" ]; then
        ssh-keygen -lf - 2>/dev/null <<< "$line"
    fi
done < <(printf '%s\n' "$RAW_SCAN") | awk '{print $2}' | sed -e 's/\r$//' | sort -u)

if [ -z "$FP_LIST" ]; then
    error "No host key fingerprints found from ssh-keyscan output."
fi

while read -r fp; do
    if [ -n "$fp" ]; then
        info "Found host key fingerprint: $fp"
    fi
done < <(echo "$FP_LIST")

while read -r fp; do
    if [ -n "$fp" ]; then
        CURL_CMD+=(--data-urlencode "fingerprint=$fp")
    fi
done < <(echo "$FP_LIST")

CA_TMP=$(mktemp /tmp/sshteam_ca.XXXXXX)
HTTP_CODE=$("${CURL_CMD[@]}" -o "$CA_TMP" -w "%{http_code}") || true
if [ "$HTTP_CODE" != "200" ]; then
    rm -f "$CA_TMP"
    error "Join request failed with HTTP $HTTP_CODE for $JOIN_URL"
fi
CA_DATA=$(cat "$CA_TMP")
rm -f "$CA_TMP"
[ -z "$CA_DATA" ] && error "No CA data returned from $JOIN_URL."

# --- 6. Configure SSH ---
info "Writing CA keys to $CA_PUB_PATH..."
echo "$CA_DATA" | sudo tee "$CA_PUB_PATH" > /dev/null
sudo chmod 644 "$CA_PUB_PATH"

if ! grep -q "TrustedUserCAKeys $CA_PUB_PATH" "$SSHD_CONFIG"; then
    info "Updating $SSHD_CONFIG with TrustedUserCAKeys..."
    echo -e "\n# SSHTEAM CA TRUST\nTrustedUserCAKeys $CA_PUB_PATH" | sudo tee -a "$SSHD_CONFIG" > /dev/null
else
    info "SaaS CA trust already configured in $SSHD_CONFIG."
fi

info "Restarting SSH daemon..."
if [[ "$OSTYPE" == "darwin"* ]]; then
    sudo launchctl unload /System/Library/LaunchDaemons/ssh.plist 2>/dev/null
    sudo launchctl load -w /System/Library/LaunchDaemons/ssh.plist
else
    if command -v systemctl >/dev/null 2>&1; then
        sudo systemctl restart ssh || sudo systemctl restart sshd
    else
        sudo service ssh restart || sudo service sshd restart
    fi
fi

ACCESS_TOKEN=""
info "Server joined and SSH configured."
