#!/usr/bin/env bash
# Self-extracting TDX guest launcher for {{CRATE_NAME}}
# Generated by mosaik TDX builder — contains vmlinuz + initramfs + ovmf.fd
#
# Usage:
#   sudo ./{{CRATE_NAME}}-run-qemu.sh
#
# Environment variables:
#   MEMORY    — guest RAM (default: {{MEMORY}})
#   CPUS      — guest vCPUs (default: {{CPUS}})
#   SSH_PORT  — host port forwarded to guest :22 (default: auto-assigned free port)
#   VSOCK_CID — guest vsock context ID (default: auto-assigned random CID)
#   OVMF      — override ovmf.fd path
#   KERNEL    — override vmlinuz path
#   INITRD    — override initramfs path
#   KEEP      — set to 1 to keep extracted files after exit

set -euo pipefail

MARKER="__PAYLOAD_BELOW__"
MEMORY="${MEMORY:-{{MEMORY}}}"
CPUS="${CPUS:-{{CPUS}}}"
VSOCK_CID="${VSOCK_CID:-auto}"

# Find a free TCP port for SSH forwarding.
find_free_port() {
    python3 -c "import socket; s=socket.socket(); s.bind(('',0)); print(s.getsockname()[1]); s.close()" 2>/dev/null \
    || perl -e 'use IO::Socket::INET; my $s=IO::Socket::INET->new(Listen=>1,LocalAddr=>"127.0.0.1",LocalPort=>0); print $s->sockport(),"\n"; $s->close();' 2>/dev/null \
    || { echo "ERROR: cannot find free port (need python3 or perl)" >&2; exit 1; }
}

# Pick a unique vsock CID (3–0xFFFFFFFF).
pick_vsock_cid() {
    local cid
    cid=$(shuf -i 3-2147483647 -n1 2>/dev/null || awk 'BEGIN{srand(); printf "%d\n", 3+int(rand()*2147483644)}')
    echo "$cid"
}

SSH_PORT="${SSH_PORT:-$(find_free_port)}"
if [ "$VSOCK_CID" = "auto" ]; then
    VSOCK_CID=$(pick_vsock_cid)
fi

WORK_DIR=$(mktemp -d /tmp/{{CRATE_NAME}}-tdx.XXXXXX)
cleanup() {
    if [ "${KEEP:-0}" = "1" ]; then
        echo "Extracted files kept at: $WORK_DIR"
    else
        rm -rf "$WORK_DIR"
    fi
}
trap cleanup EXIT

echo "Extracting TDX guest image..."
PAYLOAD_LINE=$(grep -an "^$MARKER" "$0" | tail -1 | cut -d: -f1)
PAYLOAD_START=$((PAYLOAD_LINE + 1))
tail -n +$PAYLOAD_START "$0" | tar xzf - -C "$WORK_DIR"

OVMF="${OVMF:-$WORK_DIR/ovmf.fd}"
KERNEL="${KERNEL:-$WORK_DIR/vmlinuz}"
INITRD="${INITRD:-$WORK_DIR/initramfs.cpio.gz}"

if [ ! -f "$OVMF" ]; then
    OVMF="/usr/share/ovmf/OVMF.fd"
fi

for f in "$OVMF" "$KERNEL" "$INITRD"; do
    if [ ! -f "$f" ]; then
        echo "ERROR: Missing file: $f" >&2
        exit 1
    fi
done

TDX_OBJECT='{"qom-type":"tdx-guest","id":"tdx0","sept-ve-disable":true,"quote-generation-socket":{"type":"vsock","cid":"2","port":"4050"}}'
CMDLINE="console=ttyS0 ip=dhcp"

echo "=== Launching TDX Guest: {{CRATE_NAME}} ==="
echo "  OVMF:     $OVMF"
echo "  Kernel:   $KERNEL"
echo "  Initrd:   $INITRD"
echo "  Memory:   $MEMORY"
echo "  CPUs:     $CPUS"
echo "  SSH:      localhost:$SSH_PORT"
echo "  CID:     $VSOCK_CID"
echo ""

exec qemu-system-x86_64 \
    -accel kvm \
    -cpu host,pmu=off \
    -smp "$CPUS" \
    -m "$MEMORY" \
    \
    -object "$TDX_OBJECT" \
    -machine q35,kernel-irqchip=split,confidential-guest-support=tdx0 \
    \
    -bios "$OVMF" \
    -kernel "$KERNEL" \
    -initrd "$INITRD" \
    -append "$CMDLINE" \
    \
    -netdev user,id=net0,{{SSH_FWD}} \
    -device virtio-net-pci,netdev=net0 \
    -device vhost-vsock-pci,guest-cid="$VSOCK_CID" \
    \
    -nographic \
    -nodefaults \
    -serial mon:stdio

__PAYLOAD_BELOW__