#!/usr/bin/env bash
set -euo pipefail

usage() {
  cat <<EOF
Usage: lab-addkey <user> <keyfile|key-string>

Appends an SSH public key to both:
  - the host user's ~/.ssh/authorized_keys
  - the container's /root/.ssh/authorized_keys

Examples:
  lab-addkey alice /tmp/alice-laptop.pub
  lab-addkey alice "ssh-ed25519 AAAA... user@host"
EOF
  exit 1
}

[[ $# -lt 2 ]] && usage

USER="$1"
KEY_INPUT="$2"

# resolve key content
if [[ -f "$KEY_INPUT" ]]; then
  KEY=$(cat "$KEY_INPUT")
else
  KEY="$KEY_INPUT"
fi

# basic sanity check
if [[ ! "$KEY" =~ ^ssh- ]] && [[ ! "$KEY" =~ ^ecdsa- ]]; then
  echo "ERROR: doesn't look like a valid SSH public key" >&2
  exit 1
fi

# resolve container
if [[ -f "/home/${USER}/.lxc-container" ]]; then
  CONTAINER=$(cat "/home/${USER}/.lxc-container")
else
  echo "ERROR: no container found for user '${USER}'" >&2
  exit 1
fi

# --- add to host user ---
HOST_AUTHKEYS="/home/${USER}/.ssh/authorized_keys"
mkdir -p "/home/${USER}/.ssh"
if grep -qF "$KEY" "$HOST_AUTHKEYS" 2>/dev/null; then
  echo "Key already present on host for ${USER}"
else
  echo "$KEY" >> "$HOST_AUTHKEYS"
  chown "${USER}:" "$HOST_AUTHKEYS"
  chmod 600 "$HOST_AUTHKEYS"
  echo "Added key to host: ${HOST_AUTHKEYS}"
fi

# --- add to container ---
STATE=$(lxc-info -n "$CONTAINER" -sH 2>/dev/null || true)
if [[ "$STATE" == "RUNNING" ]]; then
  echo "$KEY" | lxc-attach --clear-env -n "$CONTAINER" -- /bin/bash -c "
    export PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
    mkdir -p /root/.ssh && chmod 700 /root/.ssh
    KEY=\$(cat)
    if grep -qF \"\$KEY\" /root/.ssh/authorized_keys 2>/dev/null; then
      echo 'Key already present in container'
    else
      echo \"\$KEY\" >> /root/.ssh/authorized_keys
      chmod 600 /root/.ssh/authorized_keys
      echo 'Added key to container'
    fi
  "
else
  # container stopped — write directly into rootfs
  ROOTFS="/var/lib/lxc/${CONTAINER}/rootfs"
  CONT_AUTHKEYS="${ROOTFS}/root/.ssh/authorized_keys"
  mkdir -p "${ROOTFS}/root/.ssh"
  chmod 700 "${ROOTFS}/root/.ssh"
  if grep -qF "$KEY" "$CONT_AUTHKEYS" 2>/dev/null; then
    echo "Key already present in container rootfs"
  else
    echo "$KEY" >> "$CONT_AUTHKEYS"
    chmod 600 "$CONT_AUTHKEYS"
    echo "Added key to container rootfs: ${CONT_AUTHKEYS}"
  fi
fi

echo "Done."