# Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

CGPT_PY="${BUILD_LIBRARY_DIR}/cgpt.py"
PARTITION_SCRIPT_PATH="usr/sbin/write_gpt.sh"
DISK_LAYOUT_PATH=

cgpt_py() {
  if [[ -n "${FLAGS_adjust_part-}" ]]; then
    set -- --adjust_part "${FLAGS_adjust_part}" "$@"
    if [[ ! -t 0 ]]; then
      warn "The --adjust_part flag was passed." \
           "This option must ONLY be used interactively. If" \
           "you need to pass a size from another script, you're" \
           "doing it wrong and should be using a disk layout type."
    fi
  fi
  "${CGPT_PY}" "$@"
}

get_disk_layout_path() {
  if [[ -n ${DISK_LAYOUT_PATH} ]]; then
    return 0
  fi
  DISK_LAYOUT_PATH="${BUILD_LIBRARY_DIR}/legacy_disk_layout.json"
  local partition_script_path=$(tempfile)
  for overlay in $(cros_list_overlays --board "$BOARD"); do
    local disk_layout="${overlay}/scripts/disk_layout.json"
    if [[ -e ${disk_layout} ]]; then
      DISK_LAYOUT_PATH=${disk_layout}
    fi
  done
}

write_partition_script() {
  local image_type=$1
  local partition_script_path=$2
  get_disk_layout_path

  local temp_script_file=$(mktemp)

  sudo mkdir -p "$(dirname "${partition_script_path}")"

  cgpt_py write "${image_type}" "${DISK_LAYOUT_PATH}" \
          "${temp_script_file}"
  sudo mv "${temp_script_file}" "${partition_script_path}"
  sudo chmod a+r "${partition_script_path}"
}

run_partition_script() {
  local outdev=$1
  local root_fs_img=$2

  local pmbr_img
  case ${ARCH} in
  arm)
    pmbr_img=/dev/zero
    ;;
  amd64|x86)
    pmbr_img=$(readlink -f /usr/share/syslinux/gptmbr.bin)
    ;;
  *)
    error "Unknown architecture: $ARCH"
    return 1
    ;;
  esac

  sudo mount -o loop "${root_fs_img}" "${root_fs_dir}"
  . "${root_fs_dir}/${PARTITION_SCRIPT_PATH}"
  write_partition_table "${outdev}" "${pmbr_img}"
  sudo umount "${root_fs_dir}"
}

get_fs_block_size() {
  get_disk_layout_path

  cgpt_py readfsblocksize "${DISK_LAYOUT_PATH}"
}

get_block_size() {
  get_disk_layout_path

  cgpt_py readblocksize "${DISK_LAYOUT_PATH}"
}

get_partition_size() {
  local image_type=$1
  local part_id=$2
  get_disk_layout_path

  cgpt_py readpartsize "${image_type}" "${DISK_LAYOUT_PATH}" ${part_id}
}

get_filesystem_size() {
  local image_type=$1
  local part_id=$2
  get_disk_layout_path

  cgpt_py readfssize "${image_type}" "${DISK_LAYOUT_PATH}" ${part_id}
}

get_label() {
  local image_type=$1
  local part_id=$2
  get_disk_layout_path

  cgpt_py readlabel "${image_type}" "${DISK_LAYOUT_PATH}" ${part_id}
}

check_valid_layout() {
  local image_type=$1
  get_disk_layout_path

  cgpt_py parseonly "${image_type}" "${DISK_LAYOUT_PATH}" > /dev/null
}

get_disk_layout_type() {
  DISK_LAYOUT_TYPE="base"
  if should_build_image ${CHROMEOS_FACTORY_INSTALL_SHIM_NAME}; then
    DISK_LAYOUT_TYPE="factory_install"
  fi
}

emit_gpt_scripts() {
  local image="$1"
  local dir="$2"

  local pack="${dir}/pack_partitions.sh"
  local unpack="${dir}/unpack_partitions.sh"
  local mount="${dir}/mount_image.sh"
  local umount="${dir}/umount_image.sh"

  local start size part x

  local default

  # Write out the header for the script.
  local gpt_layout=$(${GPT} show "${image}" | sed -e 's/^/# /')
  for x in "${unpack}" "${pack}" "${mount}" "${umount}"; do
    cat >"${x}" <<\EOF
#!/bin/bash -eu
# File automatically generated. Do not edit.

usage() {
  local ret=0
  if [[ $# -gt 0 ]]; then
    # Write to stderr on errors.
    exec 1>&2
    echo "ERROR: $*"
    echo
    ret=1
  fi
  echo "Usage: $0 [image] [part]"
  echo "Example: $0 chromiumos_image.bin"
  exit ${ret}
}

TARGET=${1:-}
PART=${2:-}
case ${TARGET} in
-h|--help)
  usage
  ;;
"")
  for TARGET in chromiumos_{,base_}image.bin ""; do
    if [[ -e ${TARGET} ]]; then
      echo "autodetected image: ${TARGET}"
      break
    fi
  done
  if [[ -z ${TARGET} ]]; then
    usage "could not autodetect an image"
  fi
  ;;
*)
  if [[ ! -e ${TARGET} ]]; then
    usage "image does not exist: ${TARGET}"
  fi
esac

set -x
EOF
    echo "${gpt_layout}" >> "${x}"
  done

  # Read each partition and generate code for it.
  while read start size part x; do
    local file="part_${part}"
    local dir="dir_${part}"
    local target='"${TARGET}"'
    local dd_args="bs=512 count=${size}"
    local start_b=$(( start * 512 ))
    local size_b=$(( size * 512 ))
    local label=$(${GPT} show "${image}" -i ${part} -l)

    for x in "${unpack}" "${pack}" "${mount}" "${umount}"; do
      cat <<EOF >> "${x}"
case \${PART:-${part}} in
${part}|"${label}")
EOF
    done

    cat <<EOF >> "${unpack}"
dd if=${target} of=${file} ${dd_args} skip=${start}
ln -sfT ${file} "${file}_${label}"
EOF
    cat <<EOF >> "${pack}"
dd if=${file} of=${target} ${dd_args} seek=${start} conv=notrunc
EOF

    if [[ ${size} -gt 1 ]]; then
      cat <<-EOF >>"${mount}"
(
mkdir -p ${dir}
m=( sudo mount -o loop,offset=${start_b},sizelimit=${size_b} ${target} ${dir} )
if ! "\${m[@]}"; then
  if ! "\${m[@]}" -o ro; then
    rmdir ${dir}
    exit 0
  fi
fi
ln -sfT ${dir} "${dir}_${label}"
) &
EOF
      cat <<-EOF >>"${umount}"
if [[ -d ${dir} ]]; then
  (
  sudo umount ${dir} || :
  rmdir ${dir}
  rm -f "${dir}_${label}"
  ) &
fi
EOF
    fi

    for x in "${unpack}" "${pack}" "${mount}" "${umount}"; do
      echo "esac" >> "${x}"
    done
  done < <(${GPT} show -q "${image}")

  echo "wait" >> "${mount}"
  echo "wait" >> "${umount}"

  chmod +x "${unpack}" "${pack}" "${mount}" "${umount}"
}

build_gpt() {
  local outdev="$1"
  local rootfs_img="$2"
  local stateful_img="$3"
  local esp_img="$4"
  local oem_img="$5"

  get_disk_layout_type
  run_partition_script "${outdev}" "${rootfs_img}"

  local sudo=
  if [ ! -w "$outdev" ] ; then
    # use sudo when writing to a block device.
    sudo=sudo
  fi

  # Now populate the partitions.
  info "Copying stateful partition..."
  $sudo dd if="$stateful_img" of="$outdev" conv=notrunc bs=512 \
      seek=$(partoffset ${outdev} 1) status=none

  info "Copying rootfs..."
  $sudo dd if="$rootfs_img" of="$outdev" conv=notrunc bs=512 \
      seek=$(partoffset ${outdev} 3) status=none

  info "Copying EFI system partition..."
  $sudo dd if="$esp_img" of="$outdev" conv=notrunc bs=512 \
      seek=$(partoffset ${outdev} 12) status=none

  info "Copying OEM partition..."
  $sudo dd if="$oem_img" of="$outdev" conv=notrunc bs=512 \
      seek=$(partoffset ${outdev} 8) status=none

  # Pre-set "sucessful" bit in gpt, so we will never mark-for-death
  # a partition on an SDCard/USB stick.
  cgpt add -i 2 -S 1 "$outdev"
}

round_up_4096() {
  local blocks=$1
  local round_up=$(( blocks % 4096 ))
  if [ $round_up -ne 0 ]; then
    blocks=$(( blocks + 4096 - round_up ))
  fi
  echo $blocks
}

# Rebuild an image's partition table with new stateful size.
#  $1: source image filename
#  $2: source stateful partition image filename
#  $3: number of sectors to allocate to the new stateful partition
#  $4: destination image filename
# Used by dev/host/tests/mod_recovery_for_decryption.sh and
# mod_image_for_recovery.sh.
update_partition_table() {
  local src_img=$1              # source image
  local src_state=$2            # stateful partition image
  local dst_stateful_blocks=$3  # number of blocks in resized stateful partition
  local dst_img=$4

  rm -f "${dst_img}"

  # Make sure new stateful's size is a multiple of 4096 blocks so that
  # relocated partitions following it are not misaligned.
  dst_stateful_blocks=$(round_up_4096 $dst_stateful_blocks)
  # Calculate change in image size.
  local src_stateful_blocks=$(cgpt show -i 1 -s ${src_img})
  local delta_blocks=$(( dst_stateful_blocks - src_stateful_blocks ))
  local dst_stateful_bytes=$(( dst_stateful_blocks * 512 ))
  local src_stateful_bytes=$(( src_stateful_blocks * 512 ))
  local src_size=$(stat -c %s ${src_img})
  local dst_size=$(( src_size - src_stateful_bytes + dst_stateful_bytes ))
  truncate -s ${dst_size} ${dst_img}

  # Copy MBR, initialize GPT.
  dd if="${src_img}" of="${dst_img}" conv=notrunc bs=512 count=1 status=none
  cgpt create ${dst_img}

  # Find partition number of STATE (really should always be "1")
  local part=0
  local label=""
  while [ "${label}" != "STATE" ]; do
    part=$(( part + 1 ))
    local label=$(cgpt show -i ${part} -l ${src_img})
    local src_start=$(cgpt show -i ${part} -b ${src_img})
    if [ ${src_start} -eq 0 ]; then
      echo "Could not find 'STATE' partition" >&2
      return 1
    fi
  done
  local src_state_start=$(cgpt show -i ${part} -b ${src_img})

  # Duplicate each partition entry.
  part=0
  while :; do
    part=$(( part + 1 ))
    local src_start=$(cgpt show -i ${part} -b ${src_img})
    if [ ${src_start} -eq 0 ]; then
      # No more partitions to copy.
      break
    fi
    local dst_start=${src_start}
    # Load source partition details.
    local size=$(cgpt show -i ${part} -s ${src_img})
    local label=$(cgpt show -i ${part} -l ${src_img})
    local attr=$(cgpt show -i ${part} -A ${src_img})
    local tguid=$(cgpt show -i ${part} -t ${src_img})
    local uguid=$(cgpt show -i ${part} -u ${src_img})
    # Change size of stateful.
    if [ "${label}" = "STATE" ]; then
      size=${dst_stateful_blocks}
    fi
    # Partitions located after STATE need to have their start moved.
    if [ ${src_start} -gt ${src_state_start} ]; then
      dst_start=$(( dst_start + delta_blocks ))
    fi
    # Add this partition to the destination.
    cgpt add -i ${part} -b ${dst_start} -s ${size} -l "${label}" -A ${attr} \
             -t ${tguid} -u ${uguid} ${dst_img}
    if [ "${label}" != "STATE" ]; then
      # Copy source partition as-is.
      dd if="${src_img}" of="${dst_img}" conv=notrunc bs=512 \
        skip=${src_start} seek=${dst_start} count=${size} status=none
    else
      # Copy new stateful partition into place.
      dd if="${src_state}" of="${dst_img}" conv=notrunc bs=512 \
        seek=${dst_start} status=none
    fi
  done
  return 0
}