#!/bin/bash
# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

# Retrieved from https://cos.googlesource.com/cos/overlays/board-overlays/+/refs/heads/master/project-lakitu/dev-util/cos-dev-tools/files

#
# Convert VM instance to developer mode with mutable rootfs
#

set -o errexit
set -o pipefail

readonly ESPDEV="/dev/disk/by-partlabel/EFI-SYSTEM"
readonly GRUBCFG="/efi/boot/grub.cfg"

CONFIRMED="no"
KEEP_RO_BYDEFAULT="no"
# trigerring panic instead of reboot is preferrable
# for cos-customizer workflow, in this case customization
# will detect command failure and retry it on the instance
# with R/W-mount
TRIGGER_PANIC="no"
TRIGGER_REBOOT="no"

source /usr/share/vboot/bin/common_minimal.sh

is_dev_mode_enabled() {
  local dir
  local retval
  dir="$(mktemp -d)"

  mount -o ro "${ESPDEV}" "${dir}"
  local -r grub="${dir}/${GRUBCFG}"
  if ! grep -q "defaultA=0" "${grub}" || ! grep -q "defaultB=1" "${grub}"; then
    retval=1
  else
    retval=0
  fi

  umount "${dir}"
  rm -rf "${dir}"
  return "${retval}"
}

# This procedure may or may not trigger an immediate reboot. Cos-customizer is
# designed to re-run (and only from) the script that triggered the reboot, so
# it is important that this procedure handles re-entry gracefully.
set_dev_mode() {
  local dir
  dir="$(mktemp -d)"

  mount "${ESPDEV}" "${dir}"
  local -r grub="${dir}/${GRUBCFG}"
  # set defaultA=0 and defaultB=1
  # disable module signaturesh check
  sed -i \
    -e '/^defaultA=/s:=.*:=0:' \
    -e '/^defaultB=/s:=.*:=1:' \
    -e 's/module.sig_enforce=1//' \
    "${grub}"

  # mount rootfs as rw unless requested otherwise
  if [[ "${KEEP_RO_BYDEFAULT}" != "yes" ]]; then
    sed -i \
      -e 's/ ro / rw /' \
      "${grub}"
  fi

  sync

  local -r rootdev="$(rootdev -s)"
  blockdev --setrw "${rootdev}"
  enable_rw_mount "${rootdev}"

  umount "${dir}"
  rm -rf "${dir}"

  # Triggers an immediate reboot.
  if [[ "${TRIGGER_REBOOT}" = "yes" ]]; then
    reboot
  elif [[ "${TRIGGER_PANIC}" = "yes" ]]; then
    echo b > /proc/sysrq-trigger
  else
    echo "COS dev mode will be activated after reboot"
  fi
}

usage() {
  echo "Usage: $0 [--confirm] [--read-only] [--reboot] [--panic]"
  echo "Options:"
  echo "  --confirm   confirm operation without interactive prompt"
  echo "  --panic     force-panic after enabling dev mode"
  echo "  --read-only disable write-protect but keep rootfs read-only"
  echo "  --reboot    reboot graciously after enabling dev mode"
}

main() {
  for opt in "$@"; do
    case "${opt}" in
      --help|-h)
        usage
        exit 1
        ;;
      --confirm)
        CONFIRMED=yes
        ;;
      --panic)
        TRIGGER_PANIC=yes
        ;;
      --read-only)
        KEEP_RO_BYDEFAULT=yes
        ;;
      --reboot)
        TRIGGER_REBOOT=yes
        ;;
      *)
        echo "Unknown option: ${opt}"
        usage
        exit 1
        ;;
    esac
  done

  if ! grep -q '^ID=cos$' /etc/os-release; then
    echo "Does not appear to be a COS instance"
    echo "Refusing to set dev mode"
    exit 1
  fi

  if [[ "${EUID}" -ne 0 ]]; then
    echo "Please run as root or use sudo"
    exit
  fi

  if is_dev_mode_enabled; then
    echo "COS dev mode is already enabled"
    exit 0
  fi

  if /usr/share/cloud/is-secure-boot; then
    echo "Secure boot detected, can't convert to the dev mode"
    echo "COS dev mode is not compatible with secure boot"
    exit 1
  fi

  if [[ "${CONFIRMED}" != "yes" ]]; then
    while true; do
      read -r -p "Do you want to continue? (y/N)"
      case "${REPLY}" in
        [Yy][Ee][Ss]|[Yy]) break ;;
        [Nn][Oo]|[Nn]|"") exit 0;
      esac
    done
  fi

  set_dev_mode
}

main "$@"