#!/bin/bash

# SPDX-License-Identifier: Apache-2.0
# SPDX-License-Identifier: Unlicense

set -e
cd -- "$( dirname -- "${BASH_SOURCE[0]}" )"

echo_blue() {
  echo -e "\033[34m$1\033[0m"
}

echo_yellow() {
  echo -e "\033[0;33m$1\033[0m"
}

echo_red() {
  echo -e "\033[0;31m$1\033[0m"
}

if [[ -z "${SERVER_ADDR}" ]]; then
  echo_red "The SERVER_ADDR environment variable is needed."
  exit 1
fi

if [[ -z "${MEASUREMENT}" ]]; then
  echo_red "The MEASUREMENT environment variable is needed."
  exit 1
fi

if [[ -z "${VM_HOSTNAME}" ]]; then
  VM_HOSTNAME="detee-vm"
fi

if [[ "$LOG_LEVEL" == "DEBUG" ]]; then
  echo Environment variables:
  env | grep \
        -e SERVER_ADDR \
        -e SSH_KEY_FILE \
        -e DETEE_INSTALL_URL \
        -e DETEE_INSTALL_SHA \
        -e MEASUREMENT \
        -e VM_HOSTNAME
fi

server="$SERVER_ADDR"
ssh_pubkey_dir="${HOME}/.detee/cli/vms/ssh"
cert_dir="${HOME}/.detee/cli/vms/certs"
snp_reports_dir="${HOME}/.detee/cli/vms/snpreports"
keyfile_dir="${HOME}/.detee/cli/vms/secrets"
keyfile="${keyfile_dir}/${server}"
mkdir -p "$ssh_pubkey_dir"
mkdir -p "$cert_dir"
mkdir -p "$snp_reports_dir"
mkdir -p "$keyfile_dir"
server_crt="${cert_dir}/${server}.crt"
server_report="${snp_reports_dir}/${server}-report.bin"

mkdir -p "$(dirname "$keyfile")"
if [[ -f "$keyfile" ]]; then
  echo_yellow "Found keyfile $keyfile"
else
  echo_yellow "Creating keyfile $keyfile"
  dd if=/dev/urandom "of=$keyfile" bs=32 count=4
fi

snpguest --help > /dev/null \
  || {
  echo_yellow "Please install https://github.com/virtee/snpguest"
  exit 3
}

try_countdown=20;
echo -n Trying $server
while [[ $try_countdown -gt 0 ]]; do
  echo -n .
  curl --max-time 1 -k "https://$server" > /dev/null 2>&1 && break
  sleep 1
  ((try_countdown--))
done
echo

openssl s_client -connect "$server" </dev/null \
  | sed -n '/BEGIN CERTIFICATE/,/END CERTIFICATE/p' > "$server_crt"
openssl x509 -in "$server_crt" -noout
crt_hash=$(openssl dgst -sha3-512 < "$server_crt" | awk '{ print $2}')

echo_blue "The certificate hash is $crt_hash"

echo_blue "Getting hash from attestation report..."
ip=$(echo "$server" | cut -d ":" -f1)
port=$(echo "$server" | cut -d ":" -f2)
echo "$server" | grep -F '[' | grep -F ']' && {
  ip=$(echo "$server" | grep -oE '\[[a-f0-9\:]*\]')
  port=$(echo "$server" | grep -oE '\]:[0-9]+' | cut -d ':' -f2)
}

curl --cacert "$server_crt" \
  --resolve dtrfs-api:${port}:${ip} \
  "https://dtrfs-api:${port}/report" > "${server_report}.base64"

cat "${server_report}.base64" | basenc --base64url -d > "$server_report"
rm "${server_report}.base64" 

report_crt_hash=$( snpguest display report "$server_report" \
  | grep "Report Data" -A 4 | tail -4 | tr '\n' ' ' | sed 's/\s//g')

echo_blue "The hash in the report is $report_crt_hash"

if [[ "${crt_hash,,}" != "${report_crt_hash,,}" ]]; then
  echo The hash of the certificate does not match the hash from the report. Exiting.
  exit 2
fi

echo_blue "The hash matches!"
echo_blue "Verifying AMD signature in attestation report..."

chip_id_hash=$( snpguest display report "$server_report" \
  | grep "Chip ID:" -A 4 | tail -3 | tr '\n' ' ' | sed 's/\s//g' \
  | md5sum | awk '{ print $1 }')
tcb_hash=$(grep -e "Committed TCB" -e "Reported TCB" -A 10 "$server_report" |
  md5sum | awk '{ print $1 }')
vcek_path="${cert_dir}/${chip_id_hash}-${tcb_hash}.vcek.pem"
amd_certs_dir="${cert_dir}/amd_certs_${server}"
mkdir -p "$amd_certs_dir"

# TODO: add support for Genoa, Bergamo, Siena and Turin
# Or at least for Genoa...
[[ -f "${cert_dir}/ask-milan.pem" ]] || {
  snpguest fetch ca pem "$amd_certs_dir" milan --endorser vcek
  mv "${amd_certs_dir}/ask.pem" "${cert_dir}/ask-milan.pem" 
  mv "${amd_certs_dir}/ark.pem" "${cert_dir}/ark-milan.pem" 
}
ln -fs "${cert_dir}/ask-milan.pem" "${amd_certs_dir}/ask.pem" 
ln -fs "${cert_dir}/ark-milan.pem" "${amd_certs_dir}/ark.pem" 
[[ -f "${vcek_path}" ]] || {
  snpguest fetch vcek --processor-model milan pem "$amd_certs_dir" "$server_report" || {
    # You are probably wondering what this weird shit is doing here.
    # The AMD API for VCEK has throttling, and this scripts needs to run in parallel.
    sleep 10
    [[ -f "${vcek_path}" ]] || {
      snpguest fetch vcek --processor-model milan pem "$amd_certs_dir" "$server_report"
    }
  }
  mv "${amd_certs_dir}/vcek.pem" "${vcek_path}" 
}
ln -fs "${vcek_path}" "${amd_certs_dir}/vcek.pem" 
snpguest verify certs "$amd_certs_dir"
echo snpguest verify attestation "$amd_certs_dir" "$server_report"
snpguest verify attestation --processor-model milan "$amd_certs_dir" "$server_report"

echo_yellow "The attestation got verified based on the CA from AMD for the Milan generation!"
echo_blue "Verifying if measurement is $MEASUREMENT..."

guest_measurement=$( snpguest display report "$server_report" \
  | grep Measurement -A 3 | tail -3 | tr '\n' ' ' | sed 's/\s//g' )

echo_blue "The guests's measurement is $guest_measurement"

if [[ "${guest_measurement,,}" != "${MEASUREMENT,,}" ]]; then
  echo_red "The measurement of the server does not match."
  echo_yellow "Please use this project to get your measurement: https://github.com/virtee/sev-snp-measure"
  echo_yellow "After that, please sepcify the measurement ast the MEASUREMENT environment variable."
  exit 2
fi

echo_yellow "The measurement matched!"

echo_blue "Creating ed25519 signature..."
ed25519_signature=$(detee-cli account sign "$server_crt" || exit 1)

if [[ -n "$DETEE_INSTALL_URL" ]] && [[ -n "$DETEE_INSTALL_URL" ]]; then 
  echo_blue "Injecting installation url: ${DETEE_INSTALL_URL}"
  curl --fail-with-body \
    -H "ed25519-signature: ${ed25519_signature}" \
    --cacert "$server_crt" \
    --resolve "dtrfs-api:${port}:${ip}" \
    --data-urlencode "url=${DETEE_INSTALL_URL}" \
    -d "sha=${DETEE_INSTALL_SHA}" \
    -d "keyfile=$(cat "$keyfile" | basenc --base64url -w 0)" \
    -d "hostname=${VM_HOSTNAME}" \
    "https://dtrfs-api:${port}/install" || exit 1
else 
  echo
  echo_blue "Injecting keyfile..."
  curl --fail-with-body \
    -H "ed25519-signature: ${ed25519_signature}" \
    --cacert "$server_crt" \
    --resolve "dtrfs-api:${port}:${ip}" \
    -d "keyfile=$(cat "$keyfile" | basenc --base64url -w 0)" \
    "https://dtrfs-api:${port}/decrypt" || exit 1
fi

[[ -z "$SSH_KEY_FILE" ]] && SSH_KEY_FILE="$HOME/.ssh/id_ed25519.pub"
echo
echo_blue "Injecting SSH key from file ${SSH_KEY_FILE}"
curl --fail-with-body \
  -H "ed25519-signature: ${ed25519_signature}" \
  --cacert "$server_crt" \
  --resolve "dtrfs-api:${port}:${ip}" \
  --data-urlencode "ssh_key=$(cat ${SSH_KEY_FILE})" \
  "https://dtrfs-api:${port}/authorized_keys" || exit 1

echo
echo_blue "Downloading SSH keys from server"
curl --fail-with-body \
  -H "ed25519-signature: ${ed25519_signature}" \
  --cacert "$server_crt" \
  --resolve "dtrfs-api:${port}:${ip}" \
  "https://dtrfs-api:${port}/server_ssh_pubkeys" > "${ssh_pubkey_dir}/${server}" || {
    echo_red "Could not grab SSH pubkeys from the server. This could be a MITM attack. Error:"
    cat "${ssh_pubkey_dir}/${server}"
    echo
    echo_yellow "SSH at your own risk. IP: ${ip} Port: ${port} User: root"
    curl -X POST --cacert "$server_crt" \
      -H "ed25519-signature: ${ed25519_signature}" \
      --resolve dtrfs-api:${port}:${ip} \
      "https://dtrfs-api:${port}/switch_root" 2>/dev/null
    exit 1
  }
[[ "$port" != "22" ]] && ssh_keygen_ip="[$ip]:$port" || ssh_keygen_ip="$ip"
ssh-keygen -R "$ssh_keygen_ip" -f "${HOME}/.ssh/known_hosts2" || true
cat "${ssh_pubkey_dir}/${server}" |
  awk '{ print $1 " " $2 }' |
  xargs -I {} echo "$ssh_keygen_ip" {} >> "${HOME}/.ssh/known_hosts2"

echo
echo_blue "Starting guest OS. Give it a few seconds and try to SSH into the guest."
# this command is supposed to fail; it kills the API in the initrd
curl -X POST --cacert "$server_crt" \
  -H "ed25519-signature: ${ed25519_signature}" \
  --resolve dtrfs-api:${port}:${ip} \
  "https://dtrfs-api:${port}/switch_root" 2>/dev/null || echo Success.