#!/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

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


# Get the EFI partition of the root device. The steps are:
# 1. Get the block devices.
# 2. Filter to the block dev that has a descendant with "/" as a mount point.
# 3. Filter to the objects that have a device maj:min number.
# 4. Filter to the leaf partitions (no children).
# 5. Take the partition name from the 12th partition.
#
# This avoids several edge cases:
# 1. The /dev/disk/by-partlabel/EFI-SYSTEM symlink can be overwritten if a second disk is
#    mounted.
# 2. There may be virtual mounts for the root partition, so it might not be at the top level.
# 3. The device minor number for the 12th partition on the root device might not be 12 on
#    NVME since NVME devices share a device major number and increment the minor number for
#    each partition.
readonly ESPNAME=$(
    lsblk --json |
        jq '.blockdevices.[]' |
        jq 'select(.. | (.mountpoints? // []).[] == "/")' |
        jq '.. | select(."maj:min"? // "" != "")' |
        jq 'select(.children == null)' |
        jq -sr '.[11].name'
)
ESPDEV="/dev/${ESPNAME}"

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 [[ "${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 "$@"