#!/usr/bin/env bash

set -e

# String formatting subsystem

formatting_init() {
  endl=$'\n'
}

enquote() {
    while (( $# > 1 )); do
      printf "%q " "${1}"; shift
    done
    if (( $# == 1 )); then
      printf "%q" "${1}"; shift
    fi
}

multiline() {
  # shellcheck disable=SC1004
  echo "${1} " | awk '
    function pm(a, b, l) {
      return length(a) > l \
          && length(b) > l \
          && substr(a, 1, l+1) == substr(b, 1, l+1) \
            ? pm(a, b, l+1) : l;
    }

    !started && $0 !~ /^[ \t]*$/ {
      started=1
      match($0, /^[ \t]*/)
      prefix=substr($0, 1, RLENGTH)
    }

    started {
      print(substr($0, 1 + pm($0, prefix)));
    }
  '
}

dbg() {
  echo >&2 "$@"
}

# Cleanup subsystem (sigterm)

cleanup_init() {
  cleanup_actions=()
  trap cleanup_apply exit
}

cleanup_apply() {
  local f
  for f in "${cleanup_actions[@]}"; do
    eval "${f}"
  done
}

cleanup() {
  cleanup_actions+=("$(multiline "${1}")")
}

# Transactional execution subsystem

frag_init() {
  explain=0
  frag_transaction=()
  frag "
    #! /bin/bash
    set -e"
}

frag_apply() {
  local f
  for f in "${frag_transaction[@]}"; do
    if (( explain == 1 )); then
      dbg "${f}"
    fi
    eval "${f}"
  done
}

frag() {
  frag_transaction+=("$(multiline "${1}")")
}

frag_append() {
  local len; len="${#frag_transaction[@]}"
  frag_transaction=("${frag_transaction[@]:0:len-1}" "${frag_transaction[len-1]}${1}")
}

frag_append_esc() {
  frag_append " \\${endl}${1}"
}

# Usage documentation subsystem
usage_init() {
  usagestack=("${script}")
}

usage_snap() {
  echo "${#usagestack}"
}

usage_restore() {
  local n; n="${1}"
  dbg REST "${1}"
  usagestack=("${usagestack[@]:0:n-2}")
}


usage() {
  dbg "Usage: ${usagestack[*]}"
}

fatal() {
  dbg "FATAL: $*"
  usage
  exit 1
}

genkey() {
  usagestack+=("PRIVATE_KEYS_DIR")
  local skdir
  skdir="${1/\//}"; shift || fatal "Required positional argument: PRIVATE_KEYS_DIR"

  while (( $# > 0 )); do
    local arg; arg="$1"; shift
    case "${arg}" in
      -h | -help | --help | help) usage; return 0 ;;
      *) fatal "Unknown option ${arg}";;
    esac
  done

  if test -e "${skdir}"; then
    fatal "PRIVATE_KEYS_DIR \"${skdir}\" already exists"
  fi

  frag "
    umask 077
    mkdir -p $(enquote "${skdir}")
    wg genkey > $(enquote "${skdir}"/wgsk)
    $(enquote "${binary}") keygen \\
      private-key $(enquote "${skdir}"/pqsk) \\
      public-key  $(enquote "${skdir}"/pqpk)"
}

pubkey() {
  usagestack+=("PRIVATE_KEYS_DIR" "PUBLIC_KEYS_DIR")
  local skdir pkdir
  skdir="${1/\//}"; shift || fatal "Required positional argument: PRIVATE_KEYS_DIR"
  pkdir="${1/\//}"; shift || fatal "Required positional argument: PUBLIC_KEYS_DIR"

  while (( $# > 0 )); do
    local arg; arg="$1"; shift
    case "${arg}" in
      -h | -help | --help | help) usage; exit 0;;
      *) fatal "Unknown option ${arg}";;
    esac
  done

  if test -e "${pkdir}"; then
    fatal "PUBLIC_KEYS_DIR \"${pkdir}\" already exists"
  fi

  frag "
    mkdir -p $(enquote "${pkdir}")
    wg pubkey < $(enquote "${skdir}"/wgsk) > $(enquote "${pkdir}/wgpk")
    cp $(enquote "${skdir}"/pqpk) $(enquote "${pkdir}/pqpk")"
}

exchange() {
  usagestack+=("PRIVATE_KEYS_DIR" "[dev <device>]" "[listen <ip>:<port>]" "[peer PUBLIC_KEYS_DIR [endpoint <ip>:<port>] [persistent-keepalive <interval>] [allowed-ips <ip1>/<cidr1>[,<ip2>/<cidr2>]...]]...")
  local skdir dev lport
  dev="${project_name}0"
  skdir="${1/\//}"; shift || fatal "Required positional argument: PRIVATE_KEYS_DIR"

  while (( $# > 0 )); do
    local arg; arg="$1"; shift
    case "${arg}" in
      dev) dev="${1}"; shift || fatal "dev option requires parameter";;
      peer) set -- "peer" "$@"; break;; # Parsed down below
      listen)
        local listen; listen="${1}";
        lip="${listen%:*}";
        lport="${listen/*:/}";
        if [[ "$lip" = "$lport" ]]; then
          lip="[0::0]"
        fi
        shift;;
      -h | -help | --help | help) usage; return 0;;
      *) fatal "Unknown option ${arg}";;
    esac
  done

  if (( $# == 0 )); then
    fatal "Needs at least one peer specified"
  fi

  frag "
    # Create the Wireguard interface
    ip link add dev $(enquote "${dev}") type wireguard || true"

  cleanup "
    ip link del dev $(enquote "${dev}") || true"

  frag "
    ip link set dev $(enquote "${dev}") up"

  frag "
    # Deploy the classic wireguard private key
    wg set $(enquote "${dev}") private-key $(enquote "${skdir}/wgsk")"


  if test -n "${lport}"; then
    frag_append "listen-port $(enquote "$(( lport + 1 ))")"
  fi

  frag "
    # Launch the post quantum wireguard exchange daemon
    $(enquote "${binary}") exchange"

  if (( verbose == 1 )); then
    frag_append "verbose"
  fi

  frag_append_esc "    private-key $(enquote "${skdir}/pqsk")"
  frag_append_esc "    public-key  $(enquote "${skdir}/pqpk")"

  if test -n "${lport}"; then
    frag_append_esc "    listen $(enquote "${lip}:${lport}")"
  fi

  usagestack+=("peer" "PUBLIC_KEYS_DIR endpoint IP:PORT")

  while (( $# > 0 )); do
    shift; # Skip "peer" argument

    local peerdir ip port keepalive allowedips
    peerdir="${1/\//}"; shift || fatal "Required peer argument: PUBLIC_KEYS_DIR"

    while (( $# > 0 )); do
      local arg; arg="$1"; shift
      case "${arg}" in
        peer) set -- "peer" "$@"; break;; # Next peer
        endpoint) ip="${1%:*}"; port="${1/*:/}"; shift;;
        persistent-keepalive) keepalive="${1}"; shift;;
        allowed-ips) allowedips="${1}"; shift;;
        -h | -help | --help | help) usage; return 0;;
        *) fatal "Unknown option ${arg}";;
      esac
    done

    # Public key
    frag_append_esc "    peer public-key $(enquote "${peerdir}/pqpk")"

    # PSK
    local pskfile; pskfile="${peerdir}/psk"
    if test -f "${pskfile}"; then
      frag_append_esc "      preshared-key $(enquote "${pskfile}")"
    fi


    if test -n "${ip}"; then
      frag_append_esc "      endpoint $(enquote "${ip}:${port}")"
    fi

    frag_append_esc "      wireguard $(enquote "${dev}") $(enquote "$(cat "${peerdir}/wgpk")")"

    if test -n "${ip}"; then
      frag_append_esc "        endpoint $(enquote "${ip}:$(( port + 1 ))")"
    fi

    if test -n "${keepalive}"; then
      frag_append_esc "        persistent-keepalive $(enquote "${keepalive}")"
    fi

    if test -n "${allowedips}"; then
      frag_append_esc "        allowed-ips $(enquote "${allowedips}")"
    fi
  done
}

main() {
  formatting_init
  cleanup_init
  usage_init
  frag_init

  project_name="rosenpass"
  scriptdir="$(dirname "${script}")"
  verbose=0
  binary="$(
    find "${scriptdir}"/target/{release,debug}/"${project_name}" -printf "%T@ %p\n" 2>/dev/null \
      | sort -nr \
      | awk -v fallback="${project_name}" '
          NR == 1 { print($2) }
          END { if (NR == 0) print(fallback) }'
  )"

  # Parse command

  usagestack+=("[explain]" "[verbose]" "genkey|pubkey|exchange" "[ARGS]...")

  local cmd
  while (( $# > 0 )); do
    local arg; arg="$1"; shift
    case "${arg}" in
      genkey|pubkey|exchange) cmd="${arg}"; break;;
      explain) explain=1;;
      verbose) verbose=1;;
      -h | -help | --help | help) usage; return 0 ;;
      *) fatal "Unknown command ${arg}";;
    esac
  done

  test -n "${cmd}" || fatal "No command supplied"
  usagestack=("${script}")

  # Execute command

  usagestack+=("${cmd}")
  "${cmd}" "$@"
  usagestack=("${script}")

  # Apply transaction

  frag_apply
}

script="$0"
main "$@"