#!/usr/bin/env bash

# DEBUG
if [ ! -z "$DEBUG" ]; then
  echo "DEBUG is set"
  set -o xtrace
fi

SCRIPT_DIR=$(cd "$(dirname "$0")" && pwd)

# shellcheck source=HW5/config.sh
source "${SCRIPT_DIR}/config.sh"

function download() {
  local -r url="${1}"
  local -r sha256="${2}"

  if which nix-prefetch-url >/dev/null 2>&1; then
    nix-prefetch-url --print-path --type sha256 "${url}" "${sha256}" | tail -n 1
  else
    local -r tmpdir="$(mktemp -d)"
    cd "${tmpdir}" && local -r filename="$(curl "${url}" --remote-name -w "%{filename_effective}")"
    echo "${tmpdir}/${filename}"
  fi
}

function download_extract() {
  local -r url="${1}"
  local -r sha256="${2}"
  local -r path="${3}"

  echo "Downloading '${url}'..."
  local -r archive="$(download "${url}" "${sha256}")"

  if [ -z "$sha256" ]; then
    echo "Verify checksum of '${archive}'..."
    if [ "$(sha256sum "${archive}" | cut -d' ' -f1)" != "${sha256}" ]; then
      echo 'Error: Invalid checksum!'
      return 1
    fi
  fi

  echo "Extracting '${archive}'..."
  tar -xf "${archive}" -C "${path}" || return $?
}

function download_extract_all() {
  download_extract "${KERNEL_URL}" "${KERNEL_SHA256}" "${KERNEL_PATH}" || return $?
  download_extract "${BUSYBOX_URL}" "${BUSYBOX_SHA256}" "${BUSYBOX_PATH}" || return $?
  download_extract "${DROPBEAR_URL}" "${DROPBEAR_SHA256}" "${DROPBEAR_PATH}" || return $?
  download_extract "${GESFTPSERVER_URL}" "${GESFTPSERVER_SHA256}" "${GESFTPSERVER_PATH}" || return $?
}

function cpu_count() {
  if [ -f /proc/cpuinfo ]; then
    grep -c ^processor /proc/cpuinfo
  else
    if which sysctl >/dev/null 2>&1; then
      sysctl hw.ncpu | cut -d ' ' -f 2
    else
      echo '4'
    fi
  fi
}

# compile kernel
function compile_kernel() {
  ln -sf "${KERNEL_CONFIG_PATH}" "${KERNEL_SOURCE_PATH}/.config"
  echo "Compile kernel..."
  time ARCH="$KERNEL_ARCH" make -C "${KERNEL_SOURCE_PATH}" -j "$(cpu_count)" || return $?
  echo "Copy kernel..."
  cp -pf "${KERNEL_IMAGE}" "${ARTIFACTS_PATH}/"
}

# compile busybox
function compile_busybox() {
  ln -sf "${BUSYBOX_CONFIG_PATH}" "${BUSYBOX_SOURCE_PATH}/.config"
  echo "Compile busybox..."
  time make -C "${BUSYBOX_SOURCE_PATH}" -j "$(cpu_count)" || return $?
  echo "Copy busybox..."
  cp -pf "${BUSYBOX_BIN}" "${ARTIFACTS_PATH}/"
  echo "Copy libs and patch ld path..."
  build_binary "${BUSYBOX_BIN}" || return $?
}

# compile dropbear
function compile_dropbear() {
  ln -sf "${DROPBEAR_CONFIG_PATH}" "${DROPBEAR_SOURCE_PATH}/options.h"
  echo "Configure dropbear..."
  cd "${DROPBEAR_SOURCE_PATH}" && CC=${CROSS_COMPILE}gcc LD=${CROSS_COMPILE}ld ./configure --host=aarch64-linux-gnu --disable-lastlog --disable-syslog --disable-wtmp --disable-wtmpx --disable-utmpx || return $?
  echo "Compile dropbear..."
  time make -C "${DROPBEAR_SOURCE_PATH}" -j "$(cpu_count)" PROGRAMS="dropbear dbclient dropbearkey dropbearconvert scp" MULTI=1 || return $?
  echo "Copy dropbear..."
  cp -pf "${DROPBEAR_BIN}" "${ARTIFACTS_PATH}/"
  echo "Copy libs and patch ld path..."
  build_binary "${DROPBEAR_BIN}" || return $?
}

# compile gesftpserver
function compile_gesftpserver() {
  echo "Configure gesftpserver..."
  cd "${GESFTPSERVER_SOURCE_PATH}" && CC=${CROSS_COMPILE}gcc LD=${CROSS_COMPILE}ld CFLAGS=-pthread ./configure --host=aarch64-linux-gnu || return $?
  echo "Compile gesftpserver..."
  time make -C "${GESFTPSERVER_SOURCE_PATH}" -j "$(cpu_count)" gesftpserver || return $?
  echo "Copy gesftpserver..."
  cp -pf "${GESFTPSERVER_BIN}" "${ARTIFACTS_PATH}/"
  echo "Copy libs and patch ld path..."
  # TODO: ld.so with qemu-aarch64 does not work for gesftpserver
  #build_binary "${GESFTPSERVER_BIN}"
  patch_binary "${GESFTPSERVER_BIN}" || return $?
  copy_libs 'libc.so.6' 'libpthread.so.0' || return $?
}

# compile sysinfo
function compile_sysinfo() {
  echo "Compile sysinfo..."
  time CC=${CROSS_COMPILE}gcc LD=${CROSS_COMPILE}ld make -C "${SYSINFO_SOURCE_PATH}" -j "$(cpu_count)" || return $?
  echo "Copy sysinfo..."
  cp -pf "${SYSINFO_BIN}" "${ARTIFACTS_PATH}/"
  echo "Copy libs and patch ld path..."
  build_binary "${SYSINFO_BIN}" || return $?
}

function compile_module() {
  local -r module_dir="$1"
  local -r module="$(basename "$module_dir")"
  local -r module_test_src="${module_dir}/${module}.test.c"

  echo "Compile kernel module '$module'..."
  time KVER="$KERNEL_VERSION" make -C "${module_dir}" -j "$(cpu_count)" || return $?

  if [ -f "${module_test_src}" ]; then
    time make -C "${module_dir}" -j "$(cpu_count)" test || return $?
  fi

  echo "Copy kernel module '$module'..."
  cp -pf "${module_dir}/${module}.ko" "$MODULES_DST_DIR/"
  cp -pf "${module_dir}/${module}.ko.test" "$MODULES_DST_DIR/"

  if [ -f "${module_test_src}" ]; then
    build_binary "${module_dir}/${module}.ko.test" "$(basename "$MODULES_DST_DIR")/" || return $?
  fi
}

function modules_compile() {
  mkdir -p "$MODULES_DST_DIR"

  while IFS= read -r -d '' module_dir; do
    compile_module "$module_dir" || return $?
  done < <(find_modules)
}

function compile() {
  compile_kernel || return $?
  compile_busybox || return $?
  compile_dropbear || return $?
  compile_gesftpserver || return $?
  compile_sysinfo || return $?
}

function get_binary_linker() {
  local -r binary="$1"
  file "$binary" | grep -o -E ', interpreter (.*\.so(\.[0-9]+)?)' | cut -d' ' -f3
}

function copy_lib() {
  cp -pf "$1" "${ARTIFACTS_PATH}/lib/" || return $?
}

# copy binary libs
function copy_binary_libs() {
  local -r binary="$1"

  echo "Copy libs for '$binary'..."
  mkdir -p "${ARTIFACTS_PATH}/lib"

  local -r linker="$(get_binary_linker "$binary")"
  local -r libs="$(qemu-aarch64 "$linker" --list "$binary" | sed -e 's/[^\t].* => //' | sed -e 's/^\s//' | cut -d' ' -f1)"

  while read -r lib; do
    copy_lib "$lib" || return $?
  done <<< "$libs"
}

# copy libs
function copy_libs() {
  echo "Copy libs..."
  mkdir -p "${ARTIFACTS_PATH}/lib"
  for lib in "$@"; do
    copy_lib "$(${CROSS_COMPILE}gcc "-print-file-name=$lib")" || return $?
  done
}

# patch lib paths
function patch_binary() {
  local -r binary="$1"
  echo "Patching lib paths for '$binary'..."
  local -r linker="$(basename "$(get_binary_linker "$binary")")"
  patchelf --set-interpreter "/lib/$linker" --set-rpath '/lib' "$binary" || return $?
}

# do magic for binaries
function build_binary(){
  local -r binary="$1"
  local -r binary_dst="${ARTIFACTS_PATH}/${2}$(basename "$binary")"
  patch_binary "$binary_dst" || return $?
  copy_binary_libs "$binary" || return $?
}

# copy drobear libs
function copy_dropbear_libs() {
  local -r libs=(libnss_compat.so.2 libnss_files.so.2)
  copy_libs "${libs[@]}" || return $?
}

# create dropbear rsa key
function build_dropbear_rsa_key() {
  local -r priv_key="${ARTIFACTS_PATH}/${1:-id_dropbear}"
  local -r priv_key_openssh="${2}"
  local -r pub_key="${priv_key}.pub"

  # create private key
  if [ ! -f "${priv_key}" ]; then
    echo "Create dropbear private key..."
    qemu-aarch64 "${DROPBEAR_BIN}" dropbearkey -t rsa -s 4096 -f "${priv_key}" || return $?
    rm -f "${pub_key}"
  fi

  # convert private key to openssh
  if [ ! -z "$priv_key_openssh" ] && [ ! -f "${priv_key_openssh}" ]; then
    qemu-aarch64 "${DROPBEAR_BIN}" dropbearconvert dropbear openssh "${priv_key}" "${priv_key_openssh}" || return $?
  fi


  # create public key
  if [ ! -f "${pub_key}" ]; then
    echo "Extract dropbear public key..."
    qemu-aarch64 "${DROPBEAR_BIN}" dropbearkey -f "${priv_key}" -y | tail -n 2 | head -n 1 > "${pub_key}" || return $?
  fi
}

# build dropbear files
function build_dropbear() {
  copy_dropbear_libs || return $?
  build_dropbear_rsa_key id_dropbear "${SSH_KEY}" || return $?
  build_dropbear_rsa_key dropbear_rsa_host_key || return $?

  update_known_hosts || return $?
}

# build initrd
function build_initrd() {
  local -r name="${1:-initrd}"
  local -r root="${SCRIPT_DIR}/${name}"
  local -r output="${ARTIFACTS_PATH}/${name}.cpio"

  echo "Update libs in initrd..."
  rm -Rf "${root:?}/lib/"
  mkdir -p "${root}/lib/"
  while IFS= read -r -d '' lib; do
    ln -s "$lib" "${root}/lib/" || return $?
  done < <(find "${ARTIFACTS_PATH}/lib/" -type f -print0)
  
  echo "Build initrd..."
  cd "${root}" && find . -not -name '.keep' | cpio -L -v -o -H newc > "${output}" || return $?
}

# compile
function build() {
  build_dropbear || return $?
  build_initrd initrd || return $?
}

# run qemu
function run_qemu() {
  local -r initrd="${ARTIFACTS_PATH}/${1:-initrd.cpio}"
  local -r init="${2:-/init}"
  shift 1 && shift 1
  echo "Run qemu..."
  qemu-system-aarch64 -nographic -m 64 -M virt -cpu cortex-a53 -kernel "${KERNEL_IMAGE}" -initrd "${initrd}" -append "console=ttyAMA0,115200 loglevel=9 earlyprintk init=${init} $*" -device "virtio-net-pci,netdev=net0" -netdev "user,id=net0,net=10.4.0.0/24,dhcpstart=10.4.0.100,hostfwd=tcp::${SSH_PORT}-:22" || return $?
}

# update ssh known_hosts
function update_known_hosts() {
  echo "Update public key in known_hosts..."
  mkdir -p "${HOME}/.ssh"
  ssh-keygen -R "[${SSH_HOST}]:${SSH_PORT}"
  echo "[${SSH_HOST}]:${SSH_PORT} $(cat "${ARTIFACTS_PATH}/dropbear_rsa_host_key.pub")" >> ~/.ssh/known_hosts
}

# ssh connect
function ssh_connect() {
  local -r user="${1:-root}"
  shift 1
  echo "Connect with ssh to qemu..." >&2
  # shellcheck disable=SC2029
  ssh -p "${SSH_PORT}" -l "${user}" -i "${SSH_KEY}" "${SSH_HOST}" "$@" || return $?
}

# ssh cmd
function ssh_cmd() {
  echo "Running '$*' over ssh..." >&2
  ssh_connect root -n . /etc/profile\; "$@" || return $?
}

# scp copy
function scp_copy() {
  local -r user=root
  local -r dst="$1"
  shift 1
  echo "Copying '$*' over scp to '$dst'..."
  scp -P "${SSH_PORT}" -i "${SSH_KEY}" "$@" "root@${SSH_HOST}:${dst}" || return $?
}

function find_modules() {
  find "$MODULES_DIR" -mindepth 1 -maxdepth 1 -type d -not -name '_*' -print0
}

function modules_copy() {
  echo "Copying modules and module tests..."
  scp_copy "/lib/modules/$(ssh_cmd uname -r)/" "$MODULES_DST_DIR/"*.ko || return $?
  scp_copy "/tmp/" "$MODULES_DST_DIR/"*.ko.test || return $?
}
function modules_load() {
  echo "Load modules..."
  ssh_cmd depmod || return $?

  while IFS= read -r -d '' module_dir; do
    ssh_cmd modprobe "$(basename "$module_dir")" || return 1
  done < <(find_modules)
}
function modules_test() {
  local module
  echo "Test modules..."
  while IFS= read -r -d '' module_dir; do
    module="$(basename "$module_dir")"
    echo "Running test for kernel module '$module'"
    ssh_cmd "/tmp/${module}.ko.test" || return 1
  done < <(find_modules)
}
function modules_unload() {
  echo "Unload modules..."

  while IFS= read -r -d '' module_dir; do
    ssh_cmd rmmod -w "$(basename "$module_dir")" || return 1
  done < <(find_modules)
}

# clean untracked git files
function clean() {
  echo "Clean non git files..."
  cd "${SCRIPT_DIR}" && git clean -dfx || return $?
}

function main() {
  local -r cmd="${1:-all}"
  shift 1

  case "$cmd" in
    all )
      download_extract_all &&
      compile &&
      build ;;

    download )
      download_extract_all ;;

    compile )
      compile ;;

    build )
      build ;;

    qemu )
      run_qemu initrd.cpio /init "$@" ;;

    ssh )
      ssh_connect "$@" ;;

    ssh_cmd )
      ssh_cmd "$@" ;;

    modules )
      modules_compile &&
      modules_copy &&
      modules_load &&
      modules_test &&
      modules_unload ;;

    modules_build )
      modules_compile ;;

    modules_copy )
      modules_copy ;;

    modules_load )
      modules_load ;;

    modules_test )
      modules_test ;;

    modules_unload )
      modules_unload ;;

    clean )
      clean ;;

    * )
      echo "Error: unkown command" &&
      false ;;
  esac

  return $?
}

main "$@"