#!/bin/sh
set -eu

echo "=========================================="
echo "  docker-bastion starting"
echo "=========================================="
echo "Date:        $(date)"
echo "Hostname:    $(hostname)"
echo "OpenSSH:     $(/usr/sbin/sshd -V 2>&1 | head -1 || echo n/a)"
echo "Docker CLI:  $(docker --version 2>/dev/null || echo n/a)"
echo "=========================================="

SSH_USER="agent"
SSH_PORT="${SSH_PORT:-22}"
FORCE_COMMAND_VALUE="${FORCE_COMMAND:-}"
AUTHORIZED_KEYS_HOST="${AUTHORIZED_KEYS_HOST:-/etc/bastion/authorized_keys.host}"
AUTHORIZED_KEYS_REPO="${AUTHORIZED_KEYS_REPO:-/etc/bastion/authorized_keys.repo}"

# ---------------------------------------------------------------------------
# 1) Validate config
# ---------------------------------------------------------------------------
if [ -z "$FORCE_COMMAND_VALUE" ]; then
    echo "FATAL: FORCE_COMMAND must be set."
    echo "       e.g. FORCE_COMMAND='docker exec -it app bash'"
    echo "       or   FORCE_COMMAND='cd /workspace && ./deploy.sh'"
    exit 1
fi

# ---------------------------------------------------------------------------
# 2) Host keys — generate on first boot, persist via /etc/ssh/keys volume
# ---------------------------------------------------------------------------
echo "[1/5] Host keys..."
mkdir -p /etc/ssh/keys
chmod 700 /etc/ssh/keys
for keytype in ed25519 rsa; do
    keyfile="/etc/ssh/keys/ssh_host_${keytype}_key"
    if [ ! -f "$keyfile" ]; then
        echo "  Generating new $keytype host key"
        ssh-keygen -t "$keytype" -f "$keyfile" -N "" -q
    else
        echo "  Reusing existing $keytype host key"
    fi
    chmod 600 "$keyfile"
    [ -f "${keyfile}.pub" ] && chmod 644 "${keyfile}.pub"
done

# ---------------------------------------------------------------------------
# 3) Merge authorized_keys sources
# ---------------------------------------------------------------------------
echo "[2/5] Authorized keys..."
AUTH_FILE="/home/${SSH_USER}/.ssh/authorized_keys"
mkdir -p "$(dirname "$AUTH_FILE")"
: > "$AUTH_FILE"

added=0
for src in "$AUTHORIZED_KEYS_HOST" "$AUTHORIZED_KEYS_REPO"; do
    if [ -f "$src" ]; then
        echo "  + Merging $src"
        cat "$src" >> "$AUTH_FILE"
        # Force newline between sources (final file may not end with one).
        printf '\n' >> "$AUTH_FILE"
        added=$((added + 1))
    fi
done

if [ "$added" -eq 0 ]; then
    echo "FATAL: no authorized_keys source found."
    echo "       Mount at least one of:"
    echo "         $AUTHORIZED_KEYS_HOST   (typically ~/.ssh/authorized_keys from host)"
    echo "         $AUTHORIZED_KEYS_REPO   (typically ./docker/bastion/authorized_keys)"
    exit 1
fi

key_count=$(grep -cvE '^[[:space:]]*(#|$)' "$AUTH_FILE" 2>/dev/null || echo 0)
echo "  Loaded $key_count authorized key(s) from $added source(s)"

chown -R "${SSH_USER}:${SSH_USER}" "$(dirname "$AUTH_FILE")"
chmod 700 "$(dirname "$AUTH_FILE")"
chmod 600 "$AUTH_FILE"

# ---------------------------------------------------------------------------
# 4) ForceCommand wrapper
#
# Write the configured command to a plain file, then a small wrapper that
# exec's `sh -c "$(cat ...)"`. This way:
#   - Shell metacharacters in $FORCE_COMMAND work (&&, |, redirects, cd).
#   - The wrapper itself stays static (no escaping of user input into a
#     heredoc), and the command file is read at session start so changes
#     to $FORCE_COMMAND only need a container restart, not a rebuild.
# ---------------------------------------------------------------------------
echo "[3/5] ForceCommand..."
mkdir -p /etc/bastion
printf '%s\n' "$FORCE_COMMAND_VALUE" > /etc/bastion/force-command.cmd
chmod 0644 /etc/bastion/force-command.cmd

cat > /etc/bastion/force-command <<'WRAPPER'
#!/bin/sh
# Auto-generated by docker-bastion start-container.
# sshd invokes this script for every authenticated session.
# SSH_ORIGINAL_COMMAND is intentionally ignored — clients cannot override.
exec sh -c "$(cat /etc/bastion/force-command.cmd)"
WRAPPER
chmod 0755 /etc/bastion/force-command
echo "  $FORCE_COMMAND_VALUE"

# ---------------------------------------------------------------------------
# 5) Docker socket — if mounted, align group membership so the agent user
#    can talk to dockerd without --privileged.
# ---------------------------------------------------------------------------
echo "[4/5] Docker socket..."
if [ -S /var/run/docker.sock ]; then
    sock_gid=$(stat -c '%g' /var/run/docker.sock)
    echo "  Socket present, host gid=$sock_gid"

    grp_name=$(getent group "$sock_gid" | cut -d: -f1 || true)
    if [ -z "$grp_name" ]; then
        grp_name="dockerhost"
        addgroup -g "$sock_gid" "$grp_name" 2>/dev/null || true
        echo "  Created group $grp_name (gid=$sock_gid)"
    else
        echo "  Reusing existing group $grp_name (gid=$sock_gid)"
    fi
    addgroup "$SSH_USER" "$grp_name" 2>/dev/null || true
    echo "  Added $SSH_USER to $grp_name"
else
    echo "  WARN: /var/run/docker.sock not mounted — docker-based FORCE_COMMAND will fail."
fi

# ---------------------------------------------------------------------------
# 6) Adjust sshd_config port if non-default
# ---------------------------------------------------------------------------
if [ "$SSH_PORT" != "22" ]; then
    sed -i "s/^Port 22\$/Port ${SSH_PORT}/" /etc/ssh/sshd_config
fi

# ---------------------------------------------------------------------------
# 7) sshd config sanity check + launch
# ---------------------------------------------------------------------------
echo "[5/5] sshd config check..."
/usr/sbin/sshd -t -f /etc/ssh/sshd_config

echo "=========================================="
echo "  Listening on port ${SSH_PORT}"
echo "  User:          ${SSH_USER}"
echo "  ForceCommand:  ${FORCE_COMMAND_VALUE}"
echo "=========================================="

# -D = foreground, -e = log to stderr (so docker logs picks it up).
exec /usr/sbin/sshd -D -e