#!/bin/sh
# Joins the box it runs on to the fleet. A person runs it from the line `./fleet invite <box>` prints,
# a cloud box from cloud-init at first boot. Its arguments are the box's machine name in fleet-config.yaml, or on a
# cloud box --name-from-tag, which reads the name from the instance's Name tag and makes it the box's hostname, so
# its cloud-init never names it and a rename is a tag edit; and a Tailscale key that works once and dies within the
# hour. Every box joins the tailnet under its
# name with MagicDNS off through tailscaled-fleet, the fleet's own tailscaled with its own state and socket,
# keeps root trusting the fleet's root keys (fleet-root-keys.timer), and gets Docker with live-restore and
# Coolify's network on its fixed subnet. tailscaled-fleet runs in kernel mode with the stock tailscaled
# masked, or, beside a Tailscale the box's owner is logged in to, in user-space mode, leaving that one alone.
# `--control` adds Tailscale SSH, the fleet network, the firewall admitting only that network to the
# metadata address, and Coolify at COOLIFY_VERSION, from that release's own files with auto-update off.
# `--sync-root-keys` alone sets up the root keys, and `--move-to-tailscaled-fleet` moves the stock tailscaled's
# node to tailscaled-fleet, on a box joined before either. `--coolify` alone installs Coolify at COOLIFY_VERSION on
# a control plane joined before, or moves the Coolify there to it, returning once Coolify has restarted
# (./fleet restore coolify runs it). Every step is idempotent. POSIX sh, since the pasted line runs it as `sh` on
# a box with nothing installed yet.
set -eu

COOLIFY_VERSION=4.4.6
# Coolify's installer and upgrade script download compose files from Coolify's CDN, which serves only
# the newest release. Both run here against the pinned release's files, mirrored in the CDN's layout.
COOLIFY_RELEASE_URL=https://raw.githubusercontent.com/coollabsio/coolify/v$COOLIFY_VERSION
COOLIFY_RELEASE_MIRROR=/data/coolify/release-mirror
COOLIFY_CDN_LINE='CDN="https://cdn.coollabs.io/coolify"'
# Where the upgrade script marks each step of the restart it runs in the background, and how long to wait for it,
# as Coolify's installer waits.
COOLIFY_UPGRADE_STATUS=/data/coolify/source/.upgrade-status
COOLIFY_RESTART_SECONDS=180
FLEET_NETWORK=fleet
FLEET_SUBNET=10.255.0.0/24
# Coolify's network, the one its Traefik publishes ports on, made here before Coolify would make it from
# the pool: cloudflared on the host reaches Traefik through docker-proxy from this network's gateway,
# the one address Traefik trusts forwarded headers from (fleetcli/apply/proxy.py). deepdish's, as found.
ROOT_KEYS_SYNC=/usr/local/sbin/fleet-root-keys
# The fleet's own tailscaled: its own unit, state and socket, which fleetcli/root_ssh.py names too.
FLEET_TAILSCALE_UNIT=/etc/systemd/system/tailscaled-fleet.service
FLEET_TAILSCALE_STATE=/var/lib/tailscale-fleet
FLEET_TAILSCALE_SOCKET=/run/tailscale-fleet/tailscaled.sock
# One past Tailscale's own 41641, so the fleet's never meets a Tailscale the box runs for itself.
FLEET_TAILSCALE_PORT=41642
USERSPACE_FLAG=--tun=userspace-networking
# In user-space mode the box's own programs reach the tailnet through tailscaled-fleet's HTTP proxy here, as the
# CI runner pool reaches the registrar; fleetcli/apply/ci_runners.py names it too.
FLEET_TAILSCALE_PROXY=localhost:1055
STOCK_TAILSCALE_UNIT=tailscaled.service
STOCK_TAILSCALE_STATE=/var/lib/tailscale
COOLIFY_NETWORK=coolify
COOLIFY_SUBNET=10.0.1.0/24
COOLIFY_GATEWAY=10.0.1.1
METADATA_ADDRESS=169.254.169.254
# Where Docker takes a new network's subnet from: outside the fleet VPC (10.20.0.0/16), FLEET_SUBNET,
# COOLIFY_SUBNET and the tailnet (100.64.0.0/10). Coolify's installer keeps a daemon.json that names a pool.
DOCKER_ADDRESS_POOL=172.16.0.0/12
DOCKER_PACKAGES="docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin"

fail() {
  echo "join.sh: $1" >&2
  exit 1
}

# The instance's Name tag, read from the metadata service as IMDSv2 asks: a token first, then the tag with it.
instance_name_tag() {
  token=$(curl -fsS -X PUT "http://$METADATA_ADDRESS/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 60")
  curl -fsS -H "X-aws-ec2-metadata-token: $token" "http://$METADATA_ADDRESS/latest/meta-data/tags/instance/Name"
}

apt_install() {
  for package in "$@"; do
    if ! dpkg -s "$package" >/dev/null 2>&1; then
      apt-get update -q
      apt-get install -y -q "$@"
      return
    fi
  done
}

write_daemon_json() {
  mkdir -p /etc/docker
  [ -f /etc/docker/daemon.json ] || echo '{}' >/etc/docker/daemon.json
  jq --arg pool "$DOCKER_ADDRESS_POOL" '{
    "log-driver": "json-file",
    "log-opts": {"max-size": "10m", "max-file": "3"},
    "default-address-pools": [{"base": $pool, "size": 24}]
  } * . | .["live-restore"] = true' /etc/docker/daemon.json >/etc/docker/daemon.json.join
  mv /etc/docker/daemon.json.join /etc/docker/daemon.json
}

add_docker_repository() {
  os=$(sed -n 's/^ID=//p' /etc/os-release)
  codename=$(sed -n 's/^VERSION_CODENAME=//p' /etc/os-release)
  case $os in
    ubuntu | debian) ;;
    *) fail "Docker's apt repository covers Ubuntu and Debian; this box is $os" ;;
  esac
  install -m 0755 -d /etc/apt/keyrings
  curl -fsSL "https://download.docker.com/linux/$os/gpg" -o /etc/apt/keyrings/docker.asc
  chmod a+r /etc/apt/keyrings/docker.asc
  cat >/etc/apt/sources.list.d/docker.sources <<EOF
Types: deb
URIs: https://download.docker.com/linux/$os
Suites: $codename
Components: stable
Architectures: $(dpkg --print-architecture)
Signed-By: /etc/apt/keyrings/docker.asc
EOF
}

install_root_key_sync() {
  cat >"$ROOT_KEYS_SYNC" <<'EOF'
#!/bin/sh
# Makes root's fleet-root lines in authorized_keys the keys at ROOT_KEYS_URL, which ./fleet root-key writes:
# the current key and, once it has rotated, the one before it. Each may log in from the tailnet and the box
# itself alone, so a key that leaks opens nothing from the internet. A fetch that fails, or a list that is
# not one or two of the fleet's ed25519 keys, changes nothing. Written by join.sh; fleet-root-keys.timer
# runs it at boot and every five minutes.
set -eu
ROOT_KEYS_URL=https://fleet-root-keys-databob-labs.s3.us-east-2.amazonaws.com/fleet-root-keys.pub
FROM=100.64.0.0/10,127.0.0.1,::1
AUTHORIZED=/root/.ssh/authorized_keys

refuse() {
  echo "fleet-root-keys: $1; root's keys are unchanged" >&2
  exit 1
}

keys=$(curl -fsSL --max-time 30 "$ROOT_KEYS_URL") || refuse "cannot fetch $ROOT_KEYS_URL"
count=$(printf '%s\n' "$keys" | grep -c . || true)
if [ "$count" -lt 1 ] || [ "$count" -gt 2 ]; then
  refuse "$ROOT_KEYS_URL holds $count keys, not one or two"
fi
if printf '%s\n' "$keys" | grep -Evxq 'ssh-ed25519 [A-Za-z0-9+/]+={0,2} fleet-root'; then
  refuse "$ROOT_KEYS_URL holds a line that is not an ssh-ed25519 key named fleet-root"
fi
install -d -m 700 /root/.ssh
touch "$AUTHORIZED"
# This run's own file, so a sync the timer starts at the same moment never moves it away from under this one.
next=$(mktemp "$AUTHORIZED.XXXXXX")
trap 'rm -f "$next"' EXIT
{
  grep -v ' fleet-root$' "$AUTHORIZED" || true
  printf '%s\n' "$keys" | sed "s|^|from=\"$FROM\" |"
} >"$next"
chmod 600 "$next"
if ! cmp -s "$next" "$AUTHORIZED"; then
  mv "$next" "$AUTHORIZED"
fi
EOF
  chmod 755 "$ROOT_KEYS_SYNC"
  cat >/etc/systemd/system/fleet-root-keys.service <<EOF
[Unit]
Description=Trust the fleet's current root keys
Wants=network-online.target
After=network-online.target

[Service]
Type=oneshot
ExecStart=$ROOT_KEYS_SYNC
EOF
  cat >/etc/systemd/system/fleet-root-keys.timer <<EOF
[Unit]
Description=Trust the fleet's current root keys at boot and every five minutes

[Timer]
OnBootSec=1min
OnUnitActiveSec=5min

[Install]
WantedBy=timers.target
EOF
  systemctl daemon-reload
  systemctl enable --now fleet-root-keys.timer
  # The first sync runs now, so a box that cannot fetch the keys fails its join rather than its first apply.
  "$ROOT_KEYS_SYNC"
}

# Whether the stock tailscaled is logged in to a tailnet, or will not say; either way it is the box's owner's.
own_tailscale_taken() {
  command -v tailscale >/dev/null || return 1
  state=$(tailscale status --json 2>/dev/null | jq -r .BackendState)
  [ "$state" != NeedsLogin ] && [ "$state" != NoState ]
}

# The flags tailscaled-fleet adds to its own state, socket and port, with a leading space, or none: user-space
# with its proxy beside the owner's Tailscale, since two in kernel mode fight over one box's routes; kernel mode
# alone, so the box's own programs reach the tailnet directly. A box set up before keeps the mode it has.
fleet_tailscale_mode() {
  if [ -f "$FLEET_TAILSCALE_UNIT" ]; then
    if grep -q -- "$USERSPACE_FLAG" "$FLEET_TAILSCALE_UNIT"; then userspace_mode; fi
  elif own_tailscale_taken; then
    userspace_mode
  fi
}

userspace_mode() {
  echo " $USERSPACE_FLAG --outbound-http-proxy-listen=$FLEET_TAILSCALE_PROXY"
}

# Tailscale's binaries with the stock tailscaled masked, so no package upgrade starts it beside the fleet's.
mask_stock_tailscale() {
  if ! command -v tailscale >/dev/null; then
    curl -fsSL https://tailscale.com/install.sh | sh
  fi
  systemctl mask --now "$STOCK_TAILSCALE_UNIT"
}

install_fleet_tailscale() {
  mode=$1
  # No `tailscaled --cleanup` around it, as tailscaled.service has: beside the owner's Tailscale it would take
  # down that one's routes and firewall rules, and alone tailscaled replaces its own when it starts.
  cat >"$FLEET_TAILSCALE_UNIT" <<EOF
[Unit]
Description=The fleet's tailnet
Wants=network-pre.target
After=network-pre.target NetworkManager.service systemd-resolved.service

[Service]
ExecStart=/usr/sbin/tailscaled --state=$FLEET_TAILSCALE_STATE/tailscaled.state --statedir=$FLEET_TAILSCALE_STATE --socket=$FLEET_TAILSCALE_SOCKET --port=$FLEET_TAILSCALE_PORT$mode
RuntimeDirectory=tailscale-fleet
StateDirectory=tailscale-fleet
Restart=on-failure

[Install]
WantedBy=multi-user.target
EOF
  systemctl daemon-reload
  systemctl enable --now tailscaled-fleet.service
  waited=0
  until [ -e "$FLEET_TAILSCALE_SOCKET" ]; do
    [ "$waited" -lt 30 ] || fail "tailscaled-fleet made no socket in 30 seconds: journalctl -u tailscaled-fleet says why"
    sleep 1
    waited=$((waited + 1))
  done
}

# tailscale as the fleet's tailscaled hears it.
fleet_tailscale() {
  tailscale --socket="$FLEET_TAILSCALE_SOCKET" "$@"
}

# A box joined through the stock tailscaled keeps its node, key and tailnet address on tailscaled-fleet:
# the stock one is masked, its state copied, and the fleet's started in kernel mode.
move_to_fleet_tailscale() {
  [ ! -f "$FLEET_TAILSCALE_UNIT" ] || fail "tailscaled-fleet is set up already"
  apt_install jq
  tags=$(tailscale status --json 2>/dev/null | jq -r '.Self.Tags // [] | join(",")')
  case ",$tags," in
    *,tag:box,* | *,tag:control,*) ;;
    *) fail "the stock tailscaled is not this box on the fleet's tailnet; nothing was changed" ;;
  esac
  [ -f "$STOCK_TAILSCALE_STATE/tailscaled.state" ] || fail "no $STOCK_TAILSCALE_STATE/tailscaled.state to move"
  systemctl mask --now "$STOCK_TAILSCALE_UNIT"
  mkdir -p "$FLEET_TAILSCALE_STATE"
  cp -a "$STOCK_TAILSCALE_STATE/." "$FLEET_TAILSCALE_STATE/"
  install_fleet_tailscale ""
  fleet_tailscale status --self --peers=false
}

create_network() {
  name=$1
  shift
  if docker network inspect "$name" >/dev/null 2>&1; then
    return
  fi
  docker network create "$name" "$@"
}

install_firewall() {
  apt_install nftables
  mkdir -p /etc/fleet
  cat >/etc/fleet/firewall.nft <<EOF
# The metadata address answers the fleet network alone, so no other container on the box reaches
# the instance role. fleet-firewall.service loads this file before Docker; the flush makes a
# reload idempotent. Written by join.sh.
table inet fleet
flush table inet fleet
table inet fleet {
  chain forward {
    type filter hook forward priority -10; policy accept;
    ip daddr $METADATA_ADDRESS ip saddr != $FLEET_SUBNET drop
  }
}
EOF
  cat >/etc/systemd/system/fleet-firewall.service <<EOF
[Unit]
Description=The fleet's firewall rules
Before=docker.service

[Service]
Type=oneshot
RemainAfterExit=yes
ExecStart=/usr/sbin/nft -f /etc/fleet/firewall.nft

[Install]
WantedBy=multi-user.target
EOF
  systemctl daemon-reload
  systemctl enable fleet-firewall.service
  systemctl restart fleet-firewall.service
}

mirror_coolify_release() {
  mkdir -p "$COOLIFY_RELEASE_MIRROR"
  for file in docker-compose.yml docker-compose.prod.yml .env.production versions.json; do
    curl -fsSL "$COOLIFY_RELEASE_URL/$file" -o "$COOLIFY_RELEASE_MIRROR/$file"
  done
  for script in install.sh upgrade.sh upgrade-postgres.sh; do
    curl -fsSL "$COOLIFY_RELEASE_URL/scripts/$script" -o "$COOLIFY_RELEASE_MIRROR/$script"
  done
  for script in install.sh upgrade.sh; do
    grep -qxF "$COOLIFY_CDN_LINE" "$COOLIFY_RELEASE_MIRROR/$script" ||
      fail "Coolify $COOLIFY_VERSION's $script has no $COOLIFY_CDN_LINE line to point at the mirror"
    sed -i "s|^CDN=.*|CDN=\"file://$COOLIFY_RELEASE_MIRROR\"|" "$COOLIFY_RELEASE_MIRROR/$script"
  done
}

install_coolify() {
  if [ -f /data/coolify/source/.env ]; then
    return
  fi
  mirror_coolify_release
  # The installer names Coolify's key for this host after $USER, which cloud-init leaves unset.
  AUTOUPDATE=false USER=root HOME=/root bash "$COOLIFY_RELEASE_MIRROR/install.sh" "$COOLIFY_VERSION"
}

# Moves an installed Coolify to COOLIFY_VERSION, or restarts it there with its .env as it stands. The upgrade
# script returns before it restarts the containers, which it does in the background, writing step 6 to
# COOLIFY_UPGRADE_STATUS once they are up and removing the file 10 seconds later. Until then the Coolify that
# answers may be the one about to be replaced, so this waits for either.
upgrade_coolify() {
  mirror_coolify_release
  helper=$(jq -r '.coolify.helper.version' "$COOLIFY_RELEASE_MIRROR/versions.json")
  bash "$COOLIFY_RELEASE_MIRROR/upgrade.sh" "$COOLIFY_VERSION" "$helper"
  waited=0
  while [ -f "$COOLIFY_UPGRADE_STATUS" ] && [ "$(cut -d'|' -f1 "$COOLIFY_UPGRADE_STATUS" 2>/dev/null)" != 6 ]; do
    [ "$waited" -lt "$COOLIFY_RESTART_SECONDS" ] ||
      fail "Coolify did not restart within $COOLIFY_RESTART_SECONDS seconds: see /data/coolify/source/upgrade-*.log"
    sleep 2
    waited=$((waited + 2))
  done
}

usage="usage: join.sh <box>|--name-from-tag <key> [--control] | join.sh --sync-root-keys | join.sh --move-to-tailscaled-fleet | join.sh --coolify"
if [ "${1:-}" = "--sync-root-keys" ]; then
  [ $# -eq 1 ] || fail "$usage"
  [ "$(id -u)" -eq 0 ] || fail "run it with sudo"
  install_root_key_sync
  exit 0
fi
if [ "${1:-}" = "--coolify" ]; then
  [ $# -eq 1 ] || fail "$usage"
  [ "$(id -u)" -eq 0 ] || fail "run it with sudo"
  if [ -f /data/coolify/source/.env ]; then
    upgrade_coolify
  else
    install_coolify
  fi
  exit 0
fi
if [ "${1:-}" = "--move-to-tailscaled-fleet" ]; then
  [ $# -eq 1 ] || fail "$usage"
  [ "$(id -u)" -eq 0 ] || fail "run it with sudo"
  move_to_fleet_tailscale
  exit 0
fi
[ $# -eq 2 ] || [ $# -eq 3 ] || fail "$usage"
box=$1
key=$2
control=false
if [ $# -eq 3 ]; then
  [ "$3" = "--control" ] || fail "$usage"
  control=true
fi
name_from_tag=false
if [ "$box" = "--name-from-tag" ]; then
  name_from_tag=true
  box=$(instance_name_tag) || fail "the instance's Name tag is unreadable: is instance_metadata_tags enabled?"
fi
printf '%s' "$box" | grep -Eq '^[a-z][a-z0-9-]*$' || fail "$box is not a machine name from fleet-config.yaml"
case $key in
  tskey-auth-*) ;;
  *) fail "the key is not a Tailscale auth key (tskey-auth-...)" ;;
esac
[ "$(id -u)" -eq 0 ] || fail "run it with sudo"
if [ "$name_from_tag" = true ]; then
  hostnamectl set-hostname "$box"
fi

apt_install ca-certificates curl jq
mode=$(fleet_tailscale_mode)
if [ -z "$mode" ]; then
  mask_stock_tailscale
fi
install_fleet_tailscale "$mode"
fleet_tailscale up --auth-key="$key" --hostname="$box" --accept-dns=false --ssh="$control"
install_root_key_sync

write_daemon_json
if command -v docker >/dev/null; then
  systemctl reload docker
else
  add_docker_repository
  # shellcheck disable=SC2086
  apt_install $DOCKER_PACKAGES
fi
create_network "$COOLIFY_NETWORK" --attachable --subnet "$COOLIFY_SUBNET" --gateway "$COOLIFY_GATEWAY"

if [ "$control" = true ]; then
  create_network "$FLEET_NETWORK" --subnet "$FLEET_SUBNET"
  install_firewall
  install_coolify
fi
fleet_tailscale status --self --peers=false