#!/bin/bash
# build-ssh-extension.sh — Build a Tiny Core .tcz extension for SSH debugging
#
# Creates a self-contained SSH extension that can be loaded into KubeSolo OS
# at runtime for debugging. Uses dropbear for minimal footprint (~200 KB).
#
# Usage:
#   ./hack/build-ssh-extension.sh [--pubkey /path/to/key.pub]
#
# Output: output/ssh-debug.tcz
#
# To load on a running system:
#   1. Copy ssh-debug.tcz to /mnt/data/extensions/
#   2. Reboot, or manually: unsquashfs /mnt/data/extensions/ssh-debug.tcz -d /
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
OUTPUT_DIR="$PROJECT_ROOT/output"
BUILD_DIR="$PROJECT_ROOT/build/cache/ssh-ext"
DROPBEAR_VERSION="2024.86"

# Parse args
SSH_PUBKEY=""
for i in "$@"; do
    case $i in
        --pubkey)
            shift
            SSH_PUBKEY="$1"
            shift
            ;;
    esac
done

# Find SSH public key
if [ -z "$SSH_PUBKEY" ]; then
    for key in "$HOME/.ssh/id_ed25519.pub" "$HOME/.ssh/id_rsa.pub"; do
        if [ -f "$key" ]; then
            SSH_PUBKEY="$key"
            break
        fi
    done
fi

if [ -z "$SSH_PUBKEY" ] || [ ! -f "$SSH_PUBKEY" ]; then
    echo "ERROR: No SSH public key found."
    echo "Provide with --pubkey or ensure ~/.ssh/id_ed25519.pub exists"
    exit 1
fi

echo "==> Building SSH debug extension (.tcz)"
echo "    Public key: $SSH_PUBKEY"
echo "    Dropbear: $DROPBEAR_VERSION"

# Clean build area
rm -rf "$BUILD_DIR"
mkdir -p "$BUILD_DIR/squashfs-root"
mkdir -p "$OUTPUT_DIR"

SQUASHFS="$BUILD_DIR/squashfs-root"

# Create directory structure
mkdir -p "$SQUASHFS/usr/sbin"
mkdir -p "$SQUASHFS/usr/bin"
mkdir -p "$SQUASHFS/etc/dropbear"
mkdir -p "$SQUASHFS/root/.ssh"
mkdir -p "$SQUASHFS/usr/lib/kubesolo-os/init.d"

# Install authorized key
cp "$SSH_PUBKEY" "$SQUASHFS/root/.ssh/authorized_keys"
chmod 700 "$SQUASHFS/root/.ssh"
chmod 600 "$SQUASHFS/root/.ssh/authorized_keys"

# Download static dropbear if not cached
DROPBEAR_CACHE="$PROJECT_ROOT/build/cache/dropbear-static"
if [ ! -f "$DROPBEAR_CACHE" ]; then
    echo "==> Downloading static dropbear..."
    echo ""
    echo "NOTE: Static dropbear must be compiled separately."
    echo "For now, creating a placeholder extension structure."
    echo ""
    echo "To compile dropbear statically:"
    echo "  wget https://matt.ucc.asn.au/dropbear/releases/dropbear-${DROPBEAR_VERSION}.tar.bz2"
    echo "  tar xf dropbear-${DROPBEAR_VERSION}.tar.bz2"
    echo "  cd dropbear-${DROPBEAR_VERSION}"
    echo "  ./configure --enable-static --disable-zlib"
    echo "  make PROGRAMS='dropbear dbclient dropbearkey scp' STATIC=1"
    echo "  cp dropbear dbclient dropbearkey scp $PROJECT_ROOT/build/cache/"
    echo ""

    # Create placeholder script instead
    cat > "$SQUASHFS/usr/sbin/dropbear" << 'PLACEHOLDER'
#!/bin/sh
echo "ERROR: dropbear placeholder — compile static dropbear and rebuild extension"
exit 1
PLACEHOLDER
    chmod +x "$SQUASHFS/usr/sbin/dropbear"
else
    cp "$DROPBEAR_CACHE" "$SQUASHFS/usr/sbin/dropbear"
    chmod +x "$SQUASHFS/usr/sbin/dropbear"

    # Also copy dbclient and dropbearkey if available
    for tool in dbclient dropbearkey scp; do
        src="$PROJECT_ROOT/build/cache/${tool}-static"
        [ -f "$src" ] && cp "$src" "$SQUASHFS/usr/bin/$tool" && chmod +x "$SQUASHFS/usr/bin/$tool"
    done
fi

# Create SSH init stage
cat > "$SQUASHFS/usr/lib/kubesolo-os/init.d/85-ssh.sh" << 'EOF'
#!/bin/sh
# 85-ssh.sh — Start SSH server for debugging
# Part of ssh-debug.tcz extension

if ! command -v dropbear >/dev/null 2>&1; then
    return 0
fi

# Generate host keys on first boot
if [ ! -f /etc/dropbear/dropbear_rsa_host_key ]; then
    dropbearkey -t rsa -f /etc/dropbear/dropbear_rsa_host_key >/dev/null 2>&1
fi
if [ ! -f /etc/dropbear/dropbear_ed25519_host_key ]; then
    dropbearkey -t ed25519 -f /etc/dropbear/dropbear_ed25519_host_key >/dev/null 2>&1
fi

# Start dropbear in background
dropbear -R -p 22 2>/dev/null
echo "[kubesolo-init] SSH server (dropbear) started on port 22" >&2
EOF
chmod +x "$SQUASHFS/usr/lib/kubesolo-os/init.d/85-ssh.sh"

# Create extension info
cat > "$BUILD_DIR/ssh-debug.tcz.info" << EOF
Title:          ssh-debug.tcz
Description:    SSH debugging extension for KubeSolo OS
Version:        ${DROPBEAR_VERSION}
Author:         KubeSolo OS
Original-site:  https://github.com/portainer/kubesolo
Copying-policy: MIT
Size:           ~200KB
Extension_by:   kubesolo-os
Comments:       Provides dropbear SSH server for dev/debug access.
                NOT intended for production use.
EOF

# Build squashfs
if command -v mksquashfs >/dev/null 2>&1; then
    mksquashfs "$SQUASHFS" "$OUTPUT_DIR/ssh-debug.tcz" \
        -noappend -comp xz -b 4096
    echo ""
    echo "==> Built: $OUTPUT_DIR/ssh-debug.tcz ($(du -h "$OUTPUT_DIR/ssh-debug.tcz" | cut -f1))"
else
    echo ""
    echo "==> mksquashfs not found — extension directory prepared at:"
    echo "    $SQUASHFS"
    echo ""
    echo "    Install squashfs-tools and run:"
    echo "    mksquashfs $SQUASHFS $OUTPUT_DIR/ssh-debug.tcz -noappend -comp xz -b 4096"
fi

echo ""
echo "==> To use:"
echo "    1. Copy ssh-debug.tcz to USB drive or /mnt/data/extensions/"
echo "    2. On the target device, load with:"
echo "       unsquashfs -f -d / /mnt/data/extensions/ssh-debug.tcz"
echo "    3. Run: /usr/lib/kubesolo-os/init.d/85-ssh.sh"
echo "    4. SSH: ssh root@<device-ip>"
echo ""