#!/bin/bash
# inject-keys.sh — Post-boot SSH key and env injection for new VMs (Option B)
#
# Usage: inject-keys.sh <vm-ip> [key-dir]
#
# Expects the following files in key-dir (default: current directory):
#   - id_build          (private key for SSH tunnel)
#   - id_build.pub      (public key)
#   - build-agent.env   (filled-in env config — see build-agent.env.template)
#
# The VM must be running with temporary password auth enabled (as built by Packer).
# After injection, password auth is disabled and key-only access takes effect.
set -euo pipefail

VM_IP="${1:?Usage: inject-keys.sh <vm-ip> [key-dir]}"
KEY_DIR="${2:-.}"
BUILD_USER="build"

echo "==> Injecting keys to ${BUILD_USER}@${VM_IP} from ${KEY_DIR}"

# Verify required files exist
for f in id_build id_build.pub build-agent.env; do
    if [ ! -f "${KEY_DIR}/${f}" ]; then
        echo "ERROR: Missing ${KEY_DIR}/${f}"
        exit 1
    fi
done

# Create .ssh directory on VM
ssh -o StrictHostKeyChecking=no "${BUILD_USER}@${VM_IP}" \
    "mkdir -p ~/.ssh && chmod 700 ~/.ssh"

# Copy SSH keys
scp -o StrictHostKeyChecking=no \
    "${KEY_DIR}/id_build" "${KEY_DIR}/id_build.pub" \
    "${BUILD_USER}@${VM_IP}:~/.ssh/"

# Set correct permissions on private key
ssh -o StrictHostKeyChecking=no "${BUILD_USER}@${VM_IP}" \
    "chmod 600 ~/.ssh/id_build && chmod 644 ~/.ssh/id_build.pub"

# Add the tunnel target's host key to known_hosts (optional — agent uses
# StrictHostKeyChecking=no, but this avoids warnings in manual SSH)
echo "==> Adding workstation public key to authorized_keys"
ssh -o StrictHostKeyChecking=no "${BUILD_USER}@${VM_IP}" \
    "cat ~/.ssh/id_build.pub >> ~/.ssh/authorized_keys && chmod 600 ~/.ssh/authorized_keys"

# Copy build-agent.env to /etc (requires sudo)
echo "==> Installing build-agent.env"
scp -o StrictHostKeyChecking=no \
    "${KEY_DIR}/build-agent.env" "${BUILD_USER}@${VM_IP}:/tmp/build-agent.env"
ssh -o StrictHostKeyChecking=no "${BUILD_USER}@${VM_IP}" \
    "sudo cp /tmp/build-agent.env /etc/build-agent.env && sudo chmod 600 /etc/build-agent.env && rm /tmp/build-agent.env"

# Disable password auth (now that keys are in place)
echo "==> Disabling password authentication"
ssh -o StrictHostKeyChecking=no "${BUILD_USER}@${VM_IP}" \
    "sudo sed -i 's/^#*PasswordAuthentication.*/PasswordAuthentication no/' /etc/ssh/sshd_config && sudo systemctl restart sshd"

# Restart build-agent to pick up new env
echo "==> Restarting build-agent service"
ssh -o StrictHostKeyChecking=no -i "${KEY_DIR}/id_build" "${BUILD_USER}@${VM_IP}" \
    "sudo systemctl restart build-agent"

echo "==> Done. VM is ready. Test with: ssh -i ${KEY_DIR}/id_build ${BUILD_USER}@${VM_IP}"