#! /usr/bin/env bash

# More info at:
# - https://github.com/elitak/nixos-infect
# - https://git.selfprivacy.org/SelfPrivacy/selfprivacy-nixos-infect

: "${API_TOKEN:?API_TOKEN variable is not set}"
: "${CONFIG_URL:?CONFIG_URL variable is not set}"
: "${DNS_PROVIDER_TOKEN:?DNS_PROVIDER_TOKEN variable is not set}"
: "${DNS_PROVIDER_TYPE:?DNS_PROVIDER_TYPE variable is not set}"
: "${DOMAIN:?DOMAIN variable is not set}"
: "${ENCODED_PASSWORD:?ENCODED_PASSWORD variable is not set}"
: "${HOSTNAME:?HOSTNAME variable is not set}"
: "${LUSER:?LUSER variable is not set}"
: "${NIXOS_CONFIG_ID:?NIXOS_CONFIG_ID variable is not set}"
: "${NIX_VERSION:?NIX_VERSION variable is not set}"
: "${PROVIDER:?PROVIDER variable is not set}"
: "${STAGING_ACME:?STAGING_ACME variable is not set}"

: "${SSH_AUTHORIZED_KEY:=}"

readonly NL=$'\n'
readonly LOCAL_FLAKE_DIR="/etc/nixos"
readonly SECRETS_FILEPATH="/etc/selfprivacy/secrets.json"
readonly NIX="/root/.nix-profile/bin/nix"
readonly NIX_OPTS=(--extra-experimental-features "nix-command flakes")
DoNetConf=

genOptionalSsh() {
  [ -n "${SSH_AUTHORIZED_KEY}" ] && cat << EOF
"ssh": { "rootKeys": [ "${SSH_AUTHORIZED_KEY}" ] },
EOF
}

# TODO receive disk device from outside
determine2ndDisk() {
  case "$PROVIDER" in
      hetzner)
          echo "/dev/sdb"
          ;;
      digitalocean)
          echo "/dev/sda"
          ;;
      *)
          return 1
          ;;
  esac
}

# Merge original userdata.json with deployment specific fields and print result.
genUserdata() {
  local hashed_password diskDeviceName userdata_infect
  hashed_password="$(mkpasswd -m sha-512 "$USER_PASS")"
  diskDevice="$(determine2ndDisk)"
  diskDeviceName="${diskDevice##/dev}"

  userdata_infect=$(cat << EOF
{
    $(genOptionalSsh)
    "dns": {
        "provider": "$DNS_PROVIDER_TYPE",
        "useStagingACME": $STAGING_ACME
    },
    "server": {
        "provider": "$PROVIDER"
    },
    "domain": "$DOMAIN",
    "hashedMasterPassword": "$hashed_password",
    "hostname": "$HOSTNAME",
    "username": "$LUSER",
    "volumes": [
        {
            "device": "$diskDevice",
            "mountPoint": "/volumes/$diskDeviceName",
            "fsType": "ext4"
        }
    ],
    "modules": {
        "bitwarden": {
            "location": "$diskDeviceName"
        },
        "gitea": {
            "location": "$diskDeviceName"
        },
        "nextcloud": {
            "location": "$diskDeviceName"
        },
        "pleroma": {
            "location": "$diskDeviceName"
        },
        "simple-nixos-mailserver": {
            "location": "$diskDeviceName"
        }
    }
}
EOF
)

  jq -s '.[0] * .[1]' \
  "${1:?no userdata.json given to merge with}" <(printf "%s" "$userdata_infect")
}

genSecrets() {
  local dbpass
  dbpass="$(shuf --random-source=/dev/urandom -erz -n32 {A..Z} {a..z} {0..9} | tr -d '\n')"

  cat << EOF
{
    "api": {
        "token": "$API_TOKEN",
        "skippedMigrations": ["migrate_to_selfprivacy_channel", "mount_volume"]
    },
    "databasePassword": "$dbpass",
    "dns": {
        "apiKey": "$DNS_PROVIDER_TOKEN"
    },
    "modules": {
        "nextcloud": {
            "adminPassword": "$USER_PASS",
            "databasePassword": "$USER_PASS"
        }
    }
}
EOF
}

genHardwareConfiguration() {
  local bootcfg
  if ((isEFI)); then
    bootcfg=$(cat << EOF
  boot.loader.grub = {
    efiSupport = true;
    efiInstallAsRemovable = true;
    device = "nodev";
  };
  fileSystems."/boot" = { device = "$ESP"; fsType = "vfat"; };
EOF
)
  else
    bootcfg=$(cat << EOF
  boot.loader.grub.device = "$GRUBDEV";
EOF
)
  fi

  local currentSystem
  if ! currentSystem="$($NIX "${NIX_OPTS[@]}" eval --impure --raw --expr builtins.currentSystem)"
  then
      echo "cannot determine Nix currentSystem identifier"
      return 1
  fi

  # FIXME it's questionable whether these modules are needed at all...
  declare -a availableKernelModules=()

  [ "$PROVIDER" == "digitalocean" ] \
  && availableKernelModules+=('"ata_piix"' '"uhci_hcd"' '"xen_blkfront"')

  [ "$(uname -m)" == "x86_64" ] \
  && availableKernelModules+=('"vmw_pvscsi"')

  # TODO try nixos-generate-config first, resorting to the way below if failed
  # FIXME "nvme" is not needed for hetzner?

  cat << EOF
{ modulesPath, ... }:
{
  imports = [ (modulesPath + "/profiles/qemu-guest.nix") ];
$bootcfg
  boot.initrd.kernelModules = [ "nvme" ];
  boot.initrd.availableKernelModules = [ ${availableKernelModules[@]} ];
  fileSystems."/" = { device = "$ROOTFSDEV"; fsType = "$ROOTFSTYPE"; };
  nixpkgs.hostPlatform = "${currentSystem}";
}
EOF
}

genDeploymentConfiguration() {
  local release

  if ! release="$($NIX "${NIX_OPTS[@]}" eval --impure --raw --expr "(builtins.getFlake (builtins.toString ${LOCAL_FLAKE_DIR})).inputs.selfprivacy-nixos-config.inputs.nixpkgs.lib.trivial.release")"
  then
      echo "cannot determine NixOS release version"
      return 1
  fi

  cat << EOF
{ lib, ... }: {
  # The content below is static and belongs to this deployment only!
  # Do not copy this configuration file to another NixOS installation!

  system.stateVersion = lib.mkDefault "$release";`
`$(if [ "$DoNetConf" == "y" ]; then echo -e "$NL"; genNetworkingConf; fi)
}
EOF
}

setupConf() {
  mkdir -p ${LOCAL_FLAKE_DIR}
  if ! curl --fail "${CONFIG_URL}" \
  | tar -xz -C ${LOCAL_FLAKE_DIR} --strip-components=1 --exclude=".*"
  then
      echo "Error downloading/extracting top level flake configuration!"
      exit 1
  fi

  # generate and write hardware-configuration.nix
  if ! genHardwareConfiguration > ${LOCAL_FLAKE_DIR}/hardware-configuration.nix
  then
      echo "error generating ${LOCAL_FLAKE_DIR}/hardware-configuration.nix"
      exit 1
  fi

  # generate and write deployment.nix
  if ! genDeploymentConfiguration > ${LOCAL_FLAKE_DIR}/deployment.nix
  then
      echo "error generating ${LOCAL_FLAKE_DIR}/deployment.nix"
      exit 1
  fi

  # generate infected userdata based on original
  local userdataInfected
  userdataInfected="$(genUserdata ${LOCAL_FLAKE_DIR}/userdata.json)"
  printf "%s" "$userdataInfected" > ${LOCAL_FLAKE_DIR}/userdata.json

  # generate and write secrets
  local secrets
  secrets="$(genSecrets)"
  install -m0600 <(printf "%s" "$secrets") -DT ${SECRETS_FILEPATH}
}

# shellcheck disable=SC2207
genNetworkingConf() {
  # XXX It'd be better if we used procfs for all this...

  local IFS=$'\n'
  local eth0_name eth0_ip4s eth0_ip6s gateway gateway6 ether0 eth1_name
  local interfaces1 extraRules1 predictable_inames

  eth0_name="$(ip address show | grep '^2:' | awk -F': ' '{print $2}')"
  eth0_ip4s=($(ip address show dev "$eth0_name" | grep 'inet ' | sed -r 's|.*inet ([0-9.]+)/([0-9]+).*|{ address="\1"; prefixLength=\2; }|'))
  eth0_ip6s=($(ip address show dev "$eth0_name" | grep 'inet6 ' | sed -r 's|.*inet6 ([0-9a-f:]+)/([0-9]+).*|{ address="\1"; prefixLength=\2; }|')) || true
  gateway="$(ip route show dev "$eth0_name" | grep default | sed -r 's|default via ([0-9.]+).*|\1|')"
  gateway6="$(ip -6 route show dev "$eth0_name" | grep default | sed -r 's|default via ([0-9a-f:]+).*|\1|')" || true
  ether0="$(ip address show dev "$eth0_name" | grep link/ether | sed -r 's|.*link/ether ([0-9a-f:]+) .*|\1|')"

  eth1_name="$(ip address show | grep '^3:' | awk -F': ' '{print $2}')" || true
  if [ -n "$eth1_name" ]; then
      local eth1_ip4s eth1_ip6s ether1
      eth1_ip4s="$(ip address show dev "$eth1_name" | grep 'inet ' | sed -r 's|.*inet ([0-9.]+)/([0-9]+).*|{ address="\1"; prefixLength=\2; }|')"
      eth1_ip6s="$(ip address show dev "$eth1_name" | grep 'inet6 ' | sed -r 's|.*inet6 ([0-9a-f:]+)/([0-9]+).*|{ address="\1"; prefixLength=\2; }|')" || true
      ether1="$(ip address show dev "$eth1_name" | grep link/ether | sed -r 's|.*link/ether ([0-9a-f:]+) .*|\1|')"
      interfaces1=$(cat << EOF
      $eth1_name = {
        ipv4.addresses = [$(for a in "${eth1_ip4s[@]}"; do echo -n "
          $a"; done)
        ];
        ipv6.addresses = [$(for a in "${eth1_ip6s[@]}"; do echo -n "
          $a"; done)
        ];
      };
EOF
)
      extraRules1="ATTR{address}==\"${ether1}\", NAME=\"${eth1_name}\""
  else
      interfaces1=""
      extraRules1=""
  fi

  if [[ "$eth0_name" = eth* ]]; then
      predictable_inames="usePredictableInterfaceNames = lib.mkForce false;"
  else
      predictable_inames="usePredictableInterfaceNames = lib.mkForce true;"
  fi

  local defaultGateway6=${gateway6:+defaultGateway6 = \{ address = "${gateway6}"; interface = "${eth0_name}"; \};}
  local ipv6routes=${gateway6:+ipv6.routes = \[ \{ address = "${gateway6}"; prefixLength = 128; \} \];}
  cat << EOF
  # Networking configuration was populated by nixos-infect with the networking
  # details gathered from the running system.
  networking = {
    defaultGateway = "${gateway}";`
    `${defaultGateway6:+
    defaultGateway6}
    dhcpcd.enable = false;
    $predictable_inames
    interfaces = {
      $eth0_name = {
        ipv4.addresses = [$(for a in "${eth0_ip4s[@]}"; do echo -n "
          $a"; done)
        ];
        ipv6.addresses = [$(for a in "${eth0_ip6s[@]}"; do echo -n "
          $a"; done)
        ];
        ipv4.routes = [ { address = "${gateway}"; prefixLength = 32; } ];`
        `${ipv6routes:+
        $ipv6routes}
      };`
`${interfaces1:+
$interfaces1}
    };
  };
  services.udev.extraRules = ''
    ATTR{address}=="${ether0}", NAME="${eth0_name}"`
    `${extraRules1:+
    $extraRules1}
  '';
EOF
}

makeSwap() {
  # TODO check currently available swapspace first
  swapFile=$(mktemp /tmp/nixos-infect.XXXXX.swp)
  dd if=/dev/zero "of=$swapFile" bs=1M count=$((1*1024))
  chmod 0600 "$swapFile"
  mkswap "$swapFile"
  swapon -v "$swapFile"
}

removeSwap() {
  swapoff -a
  rm -vf /tmp/nixos-infect.*.swp
}

findESP() {
  local esp
  for d in /boot/EFI /boot/efi /boot; do
    [[ ! -d "$d" ]] && continue
    [[ "$d" == "$(df "$d" --output=target | sed 1d)" ]] \
      && esp="$(df "$d" --output=source | sed 1d)" \
      && break
  done
  [[ -z "$esp" ]] && { echo "ERROR: No ESP mount point found"; return 1; }
  for uuid in /dev/disk/by-uuid/*; do
    [[ $(readlink -f "$uuid") == "$esp" ]] && echo "$uuid" && return 0
  done
}

prepareEnv() {
  if ! USER_PASS="$(base64 -d <<<"$ENCODED_PASSWORD")"; then
      echo "Error decoding ENCODED_PASSWORD from Base64!"
      exit 1
  fi
  readonly USER_PASS

  isEFI=0
  [ -d /sys/firmware/efi ] && isEFI=1

  if ((isEFI)); then
    ESP="$(findESP)"
  else
    for GRUBDEV in /dev/vda /dev/sda /dev/nvme0n1; do
        [[ -e $GRUBDEV ]] && break;
    done
  fi

  # Retrieve root fs block device
  #                   (get root mount)  (get partition or logical volume)
  ROOTFSDEV=$(mount | grep "on / type" | awk '{print $1;}')
  ROOTFSTYPE=$(df "$ROOTFSDEV" --output=fstype | sed 1d)

  # DigitalOcean doesn't seem to set USER while running user data
  export USER="root"
  export HOME="/root"

  # Nix installer tries to use sudo regardless of whether we're already uid 0
  #which sudo || { sudo() { eval "$@"; }; export -f sudo; }
  # shellcheck disable=SC2174
  mkdir -p -m 0755 /nix
}

fakeCurlUsingWget() {
  # Use adapted wget if curl is missing
  which wget && { \
    curl() {
      eval "wget $(
        (local isStdout=1
        for arg in "$@"; do
          case "$arg" in
            "-o")
              echo "-O";
              isStdout=0
              ;;
            "-O")
              isStdout=0
              ;;
            "-L")
              ;;
            *)
              echo "$arg"
              ;;
          esac
        done;
        [[ $isStdout -eq 1 ]] && echo "-O-"
        )| tr '\n' ' '
      )"
    }; export -f curl; }
}

req() {
  type "$1" > /dev/null 2>&1 || which "$1" > /dev/null 2>&1
}

checkEnv() {
  [[ "$(whoami)" == "root" ]] || { echo "ERROR: Must run as root"; return 1; }

  # Perform some easy fixups before checking
  # TODO prevent multiple calls to apt-get update
  (which dnf && dnf install -y perl-Digest-SHA) || true # Fedora 24
  which xzcat || (which yum && yum install -y xz-utils) \
              || (which apt-get && apt-get update && apt-get install -y xz-utils) \
              || true
  which curl  || fakeCurlUsingWget \
              || (which apt-get && apt-get update && apt-get install -y curl) \
              || true

  req curl || req wget || { echo "ERROR: Missing both curl and wget";  return 1; }
  req xzcat            || { echo "ERROR: Missing xzcat";               return 1; }
  req awk              || { echo "ERROR: Missing awk";                 return 1; }
  req cut || req df    || { echo "ERROR: Missing coreutils (cut, df)"; return 1; }
  req mkpasswd         || { echo "ERROR: Missing mkpasswd";            return 1; }
  req shuf             || { echo "ERROR: Missing shuf";                return 1; }
}

# Download and execute the nix installer script.

installNix() {
  # install multiuser (system-wide with nix-daemon) Nix in the current system

  local nixReleaseBase='https://releases.nixos.org'
  local installURL="${nixReleaseBase}/nix/nix-${NIX_VERSION}/install"
  local shaURL="${installURL}.sha256"
  local sha tmpNixInstall

  # temporary destination for install script
  tmpNixInstall="$(mktemp -t nix-install-XXXXXXXXXX)"
  if [[ ! -f "${tmpNixInstall}" ]]; then
      echo "Failed creating a temporary file for Nix install script!"
      return 1
  fi

  echo "Downloading install script from ${installURL}..."
  if ! curl --fail "${installURL}" -o "${tmpNixInstall}" &>/dev/null; then
      echo "Failure while downloading Nix install script!"
      return 1
  fi

  if ! sha="$(curl --fail "${shaURL}")"; then
      echo "Failure while downloading Nix install script sha!"
      return 1
  fi

  echo "Validating Nix install script checksum..."
  if ! echo "${sha}  ${tmpNixInstall}" | sha256sum -c; then
      echo "Checksum validation failed!"
      return 1
  fi

  echo "Running nix installer..."
  if $SHELL "${tmpNixInstall}" \
    --daemon --no-channel-add --daemon-user-count 4; then
      echo "Nix is installed"
      rm "${tmpNixInstall}"
  else
      echo "Nix installation script failed!"
      return 1
  fi
}

infect() {
  # this is needed solely for accepting the sp-module sub-flake
  # see https://github.com/NixOS/nix/issues/3978#issuecomment-952418478
  $NIX "${NIX_OPTS[@]}" flake lock ${LOCAL_FLAKE_DIR} --update-input sp-modules

  echo "nix build the configuration flake..."
  if ! $NIX "${NIX_OPTS[@]}" build \
  --profile /nix/var/nix/profiles/system \
  ${LOCAL_FLAKE_DIR}/#nixosConfigurations."$NIXOS_CONFIG_ID".config.system.build.toplevel
  then
      echo "Failed!"
      exit 1
  fi

  # Reify resolv.conf (???)
  [[ -L /etc/resolv.conf ]] && mv -v /etc/resolv.conf /etc/resolv.conf.lnk && cat /etc/resolv.conf.lnk > /etc/resolv.conf

  # Stage the Nix coup d'état
  touch /etc/NIXOS
  echo etc/nixos            > /etc/NIXOS_LUSTRATE
  echo etc/resolv.conf     >> /etc/NIXOS_LUSTRATE
  echo ${SECRETS_FILEPATH} >> /etc/NIXOS_LUSTRATE

  rm -rf /boot.bak
  ((isEFI)) && umount "$ESP"
  mv -v /boot /boot.bak
  if ((isEFI)); then
    mkdir /boot
    mount "$ESP" /boot
    find /boot -depth ! -path /boot -exec rm -rf {} +
  fi

  echo "make configuration boot by default..."
  if ! /nix/var/nix/profiles/system/bin/switch-to-configuration boot; then
      echo "Failed!"; exit 1
  fi

  # Remove nix installed by the "install" script.
  rm -fv /nix/var/nix/profiles/default*
  /nix/var/nix/profiles/system/sw/bin/nix-collect-garbage
}

set -o errtrace
set -o nounset
set -o pipefail
set -o xtrace
shopt -s inherit_errexit
trap 'echo ${LINENO}: "$BASH_COMMAND"; exit 1' ERR

genNetworkingConf

# digitalocean requires detailed network config to be generated
[ "$PROVIDER" == "digitalocean" ] && DoNetConf="y"

apt update
apt install -y git tar curl whois jq

checkEnv
prepareEnv
makeSwap # smallest (512MB) droplet needs extra memory!
installNix
setupConf
infect
removeSwap

if [[ -z "${NO_REBOOT+x}" ]]; then
  reboot
fi