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

usage() {
  cat <<'EOF'
Usage: sudo bash setup_machronix_devtunnel.sh [server-label]

Examples:
  sudo bash setup_machronix_devtunnel.sh primary
  sudo bash setup_machronix_devtunnel.sh iotb
  sudo bash setup_machronix_devtunnel.sh iotc

The first run opens a GitHub device-login prompt. The persistent tunnel ID is
stored in /etc/machronix-devtunnel/tunnel-id and reused on later runs.
EOF
}

if [[ ${1:-} == "-h" || ${1:-} == "--help" ]]; then
  usage
  exit 0
fi

if [[ ${EUID} -ne 0 ]]; then
  echo "Run this script as root with sudo." >&2
  exit 1
fi

if [[ $(uname -m) != "x86_64" ]]; then
  echo "This script currently supports Linux x86_64 only." >&2
  exit 1
fi

label=${1:-$(hostname -s)}
label=$(printf '%s' "$label" | tr '[:upper:]' '[:lower:]' | tr -cs 'a-z0-9-' '-')
label=${label#-}
label=${label%-}
if [[ -z $label ]]; then
  echo "The server label must contain a letter or number." >&2
  exit 1
fi
label=${label:0:30}

state_dir=/etc/machronix-devtunnel
tunnel_id_file=$state_dir/tunnel-id
tunnel_port=22022
devtunnel_bin=/usr/local/bin/devtunnel

export HOME=/root
export DEBIAN_FRONTEND=noninteractive

apt-get update -qq
apt-get install -y -qq ca-certificates curl openssl socat

download_dir=$(mktemp -d)
cleanup() {
  rm -rf -- "$download_dir"
}
trap cleanup EXIT

curl --fail --silent --show-error --location \
  https://aka.ms/TunnelsCliDownload/linux-x64 \
  --output "$download_dir/devtunnel"
chmod 0700 "$download_dir/devtunnel"
"$download_dir/devtunnel" --version >/dev/null
install -m 0755 "$download_dir/devtunnel" "$devtunnel_bin"

login_status=$("$devtunnel_bin" user show --json 2>/dev/null || true)
if ! grep -Eq '"status"[[:space:]]*:[[:space:]]*"Logged in"' <<<"$login_status"; then
  echo "Authorize this server with GitHub when prompted."
  "$devtunnel_bin" user login --github --use-device-code-auth
fi

login_status=$("$devtunnel_bin" user show --json 2>/dev/null || true)
if ! grep -Eq '"status"[[:space:]]*:[[:space:]]*"Logged in"' <<<"$login_status"; then
  echo "Dev Tunnel authentication did not complete." >&2
  exit 1
fi

install -d -m 0700 "$state_dir"

if [[ -s $tunnel_id_file ]]; then
  tunnel_id=$(tr -d '[:space:]' < "$tunnel_id_file")
  if ! "$devtunnel_bin" show "$tunnel_id" >/dev/null 2>&1; then
    echo "Stored tunnel '$tunnel_id' is unavailable. Re-authenticate or remove" >&2
    echo "$tunnel_id_file only if you intentionally want a new tunnel." >&2
    exit 1
  fi
else
  tunnel_id="machronix-${label}-$(openssl rand -hex 4)"
  "$devtunnel_bin" create "$tunnel_id" \
    --description "Machronix authenticated SSH access for ${label}" \
    --expiration 30d
  umask 077
  printf '%s\n' "$tunnel_id" > "$tunnel_id_file"
fi

if ! "$devtunnel_bin" port show "$tunnel_id" --port-number "$tunnel_port" >/dev/null 2>&1; then
  "$devtunnel_bin" port create "$tunnel_id" \
    --port-number "$tunnel_port" \
    --protocol auto \
    --description "Loopback-only SSH proxy"
fi

cat > /etc/systemd/system/machronix-devtunnel-ssh-proxy.service <<EOF
[Unit]
Description=Machronix loopback SSH proxy for Microsoft Dev Tunnel
After=network.target ssh.service

[Service]
Type=simple
ExecStart=/usr/bin/socat TCP-LISTEN:${tunnel_port},bind=127.0.0.1,reuseaddr,fork TCP:127.0.0.1:22
Restart=always
RestartSec=5
NoNewPrivileges=true
PrivateTmp=true
ProtectHome=true
ProtectSystem=strict

[Install]
WantedBy=multi-user.target
EOF

cat > /etc/systemd/system/machronix-devtunnel.service <<EOF
[Unit]
Description=Machronix authenticated Microsoft Dev Tunnel
After=network-online.target machronix-devtunnel-ssh-proxy.service
Wants=network-online.target
Requires=machronix-devtunnel-ssh-proxy.service
StartLimitIntervalSec=0

[Service]
Type=simple
User=root
Environment=HOME=/root
ExecStart=${devtunnel_bin} host ${tunnel_id}
Restart=always
RestartSec=10
KillSignal=SIGINT
TimeoutStopSec=30
NoNewPrivileges=true
PrivateTmp=true
ProtectHome=false
ProtectSystem=full

[Install]
WantedBy=multi-user.target
EOF

cat > /usr/local/sbin/machronix-devtunnel-check <<'EOF'
#!/usr/bin/env bash
set -Eeuo pipefail

for unit in machronix-devtunnel-ssh-proxy.service machronix-devtunnel.service; do
  if ! systemctl is-active --quiet "$unit"; then
    systemctl restart "$unit"
  fi
done
EOF
chmod 0755 /usr/local/sbin/machronix-devtunnel-check

cat > /etc/systemd/system/machronix-devtunnel-watch.service <<'EOF'
[Unit]
Description=Check the Machronix Microsoft Dev Tunnel
After=network-online.target

[Service]
Type=oneshot
ExecStart=/usr/local/sbin/machronix-devtunnel-check
EOF

cat > /etc/systemd/system/machronix-devtunnel-watch.timer <<'EOF'
[Unit]
Description=Check the Machronix Microsoft Dev Tunnel every two minutes

[Timer]
OnBootSec=1min
OnUnitActiveSec=2min
Persistent=true

[Install]
WantedBy=timers.target
EOF

systemctl daemon-reload
systemctl enable --now \
  machronix-devtunnel-ssh-proxy.service \
  machronix-devtunnel.service \
  machronix-devtunnel-watch.timer

ready=false
for _ in $(seq 1 30); do
  if systemctl is-active --quiet machronix-devtunnel.service && \
     journalctl -u machronix-devtunnel.service --since "2 minutes ago" --no-pager \
       | grep -q "Ready to accept connections"; then
    ready=true
    break
  fi
  sleep 1
done

if [[ $ready != true ]]; then
  echo "Tunnel service did not become ready. Inspect:" >&2
  echo "  journalctl -u machronix-devtunnel.service -n 100 --no-pager" >&2
  exit 1
fi

echo
echo "Tunnel is running."
echo "Tunnel ID: $tunnel_id"
echo "Remote SSH tunnel port: $tunnel_port"
journalctl -u machronix-devtunnel.service --since "2 minutes ago" --no-pager \
  | grep "Hosting port" | tail -n 1 || true
echo
echo "Create a temporary connect-only token when access is needed:"
echo "  sudo devtunnel token $tunnel_id --scopes connect"
echo
echo "Share only the Tunnel ID and connect token. Do not share a host or management token."
echo "The client connects with:"
echo "  devtunnel connect $tunnel_id --access-token '<CONNECT_TOKEN>'"
echo "Then, in another terminal:"
echo "  ssh -p $tunnel_port root@127.0.0.1"
echo
echo "Status:"
systemctl --no-pager --full status machronix-devtunnel.service | sed -n '1,12p'