#!/bin/sh
# shellz installer. Usage:
#   curl -fsSL https://get.shellz.dev | sudo sh -s -- szt_<token>
#
# Downloads the signed shellz agent for this machine, verifies the release
# signature (Ed25519, key below) and checksum, then enrolls this server with
# the one-time join token. Requires root, systemd, curl and OpenSSL 3.
set -eu

RELEASE_PUBLIC_KEY="9YU707-WauBHpN_NkzlnLx9dswcdaCJi3EPyFDoWaQ0"
RELEASES_URL="${SHELLZ_RELEASES_URL:-https://releases.shellz.dev}"
BIN=/usr/local/bin/shellz

if [ -t 1 ]; then
	B="$(printf '\033[1m')" D="$(printf '\033[2m')" G="$(printf '\033[32m')" R="$(printf '\033[31m')" N="$(printf '\033[0m')"
else
	B="" D="" G="" R="" N=""
fi

step() { printf '  %s→%s %s\n' "$D" "$N" "$*"; }
ok() { printf '  %s✓%s %s\n' "$G" "$N" "$*"; }
die() {
	printf '  %s✗ %s%s\n' "$R" "$*" "$N" >&2
	exit 1
}

need() {
	command -v "$1" >/dev/null 2>&1 || die "$1 is required but not installed. $2"
}

fetch() {
	case "$1" in
	https://*) curl -fsSL --proto '=https' --tlsv1.2 --retry 3 --connect-timeout 15 -o "$2" "$1" ;;
	*) curl -fsSL --retry 3 --connect-timeout 15 -o "$2" "$1" ;;
	esac
}

# b64url → standard base64 with padding.
b64std() {
	s=$(printf '%s' "$1" | tr '_-' '/+')
	case $((${#s} % 4)) in
	2) s="$s==" ;;
	3) s="$s=" ;;
	esac
	printf '%s' "$s"
}

hex() { od -An -v -tx1 | tr -d ' \n'; }

# field NAME FILE: value of a one-key-per-line JSON manifest.
field() {
	sed -n 's/^[[:space:]]*"'"$1"'":[[:space:]]*"\{0,1\}\([^",]*\)"\{0,1\},\{0,1\}[[:space:]]*$/\1/p' "$2" | head -n 1
}

check_token() {
	token=$1
	case "$token" in
	szt_*) ;;
	*) die "That is not a shellz join token (it starts with szt_). Copy the full command from the shellz dashboard." ;;
	esac
	body=${token#szt_}
	case "$body" in
	*[!a-z2-7]*) die "The join token contains invalid characters. Copy the full command from the shellz dashboard." ;;
	esac
	[ ${#body} -eq 61 ] || die "The join token is incomplete. Copy the full command from the shellz dashboard."
	# base32-decode to octal escapes: secret(20) || fingerprint prefix(16) || checksum(2).
	# shellcheck disable=SC2059
	printf "$(base32_octal "$body")" >"$tmp/token"
	want=$(tail -c 2 "$tmp/token" | hex)
	got=$(head -c 36 "$tmp/token" | openssl dgst -sha256 -binary | head -c 2 | hex)
	[ "$want" = "$got" ] || die "The join token is damaged (checksum mismatch). Copy the full command from the shellz dashboard."
}

base32_octal() {
	awk -v s="$1" 'BEGIN {
		alphabet = "abcdefghijklmnopqrstuvwxyz234567"; bits = 0; value = 0
		for (i = 1; i <= length(s); i++) {
			value = value * 32 + index(alphabet, substr(s, i, 1)) - 1; bits += 5
			if (bits >= 8) { bits -= 8; byte = int(value / 2 ^ bits); value -= byte * 2 ^ bits; printf "\\%03o", byte }
		}
	}'
}

main() {
	printf '\n  %sshellz%s %s· session recording for every server%s\n\n' "$B" "$N" "$D" "$N"
	[ $# -eq 1 ] || die "Usage: curl -fsSL https://get.shellz.dev | sudo sh -s -- <join token>"
	[ "$(id -u)" -eq 0 ] || die "Please run as root: curl -fsSL https://get.shellz.dev | sudo sh -s -- <join token>"
	[ "$(uname -s)" = Linux ] || die "shellz supports Linux servers only."
	[ -d /run/systemd/system ] || die "shellz needs a server booted with systemd."
	case "$(uname -m)" in
	x86_64 | amd64) arch=amd64 ;;
	aarch64 | arm64) arch=arm64 ;;
	*) die "Unsupported CPU architecture $(uname -m); shellz supports x86_64 and arm64." ;;
	esac
	need curl "Install it (e.g. apt-get install curl) and run the command again."
	need openssl "Install OpenSSL 3 (e.g. apt-get install openssl) and run the command again."
	case "$(openssl version 2>/dev/null)" in
	"OpenSSL "[3-9].*) ;;
	*) die "OpenSSL 3 or newer is required to verify the release signature (found: $(openssl version 2>/dev/null || echo none))." ;;
	esac
	if [ -e /var/lib/shellz/host.json ]; then
		die "This server is already enrolled. Run 'sudo shellz status' to check it."
	fi

	tmp=$(mktemp -d)
	trap 'rm -rf "$tmp"' EXIT
	trap 'exit 130' INT TERM
	check_token "$1"

	step "Fetching the latest release for linux/$arch"
	fetch "$RELEASES_URL/stable/linux-$arch.json" "$tmp/manifest.json" || die "Could not download the release manifest from $RELEASES_URL."
	version=$(field version "$tmp/manifest.json")
	m_arch=$(field arch "$tmp/manifest.json")
	url=$(field url "$tmp/manifest.json")
	sha=$(field sha256 "$tmp/manifest.json")
	size=$(field size "$tmp/manifest.json")
	signature=$(field signature "$tmp/manifest.json")
	[ "$m_arch" = "$arch" ] || die "The release manifest is for $m_arch, not $arch."
	printf '%s' "$version" | grep -Eq '^[0-9]+\.[0-9]+\.[0-9]+(-[0-9A-Za-z.-]+)?$' || die "The release manifest is invalid."
	printf '%s' "$sha" | grep -Eq '^[0-9a-f]{64}$' || die "The release manifest is invalid."
	printf '%s' "$size" | grep -Eq '^[1-9][0-9]*$' || die "The release manifest is invalid."
	case "$url" in
	"$RELEASES_URL"/*) ;;
	*) die "The release manifest points outside $RELEASES_URL." ;;
	esac

	printf 'shellz:release:v1:%s\nlinux\n%s\n%s\n%s\n%s' "$version" "$arch" "$url" "$sha" "$size" >"$tmp/signed"
	printf -- '-----BEGIN PUBLIC KEY-----\nMCowBQYDK2VwAyEA%s\n-----END PUBLIC KEY-----\n' "$(b64std "$RELEASE_PUBLIC_KEY")" >"$tmp/release.pem"
	b64std "$signature" | openssl base64 -d -A >"$tmp/signature" 2>/dev/null || die "The release signature is malformed."
	[ "$(wc -c <"$tmp/signature" | tr -d ' ')" -eq 64 ] || die "The release signature is malformed."
	openssl pkeyutl -verify -pubin -inkey "$tmp/release.pem" -rawin -in "$tmp/signed" -sigfile "$tmp/signature" >/dev/null 2>&1 ||
		die "The release signature does not verify. Refusing to install."
	ok "Release $version signature verified"

	step "Downloading shellz $version"
	fetch "$url" "$tmp/shellz" || die "Could not download $url."
	[ "$(wc -c <"$tmp/shellz" | tr -d ' ')" = "$size" ] || die "The download has the wrong size. Refusing to install."
	got=$(openssl dgst -sha256 -r "$tmp/shellz" | cut -d ' ' -f 1)
	[ "$got" = "$sha" ] || die "The download checksum does not match the signed release. Refusing to install."
	ok "Checksum verified"

	mkdir -p "$(dirname "$BIN")"
	if command -v chattr >/dev/null 2>&1 && [ -e "$BIN" ]; then chattr -i "$BIN" 2>/dev/null || true; fi
	cp "$tmp/shellz" "$BIN.new"
	chmod 0755 "$BIN.new"
	mv -f "$BIN.new" "$BIN"
	[ "$("$BIN" version --short 2>/dev/null)" = "$version" ] || die "The installed binary does not run on this server."

	if ! "$BIN" enroll "$1"; then
		[ -e /var/lib/shellz/host.json ] || rm -f "$BIN"
		exit 1
	fi
	printf '\n  %sThis server is now recording.%s Watch sessions live at https://app.shellz.dev\n\n' "$B" "$N"
}

main "$@"