#!/bin/bash

# Copyright (c) 2011 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.

# Script to generate a Chromium OS update for use by the update engine.
# If a source .bin is specified, the update is assumed to be a delta update.

# Load common CrOS utilities.  Inside the chroot this file is installed in
# /usr/lib/crosutils.  This script may also be called from a zipfile, in which
# case common.sh will be in the current directory.
find_common_sh() {
  local thisdir="$(dirname "$(readlink -f "$0")")"
  local common_paths=(/usr/lib/crosutils "${thisdir}")
  local path

  SCRIPT_ROOT="${common_paths[0]}"
  for path in "${common_paths[@]}"; do
    if [[ -r "${path}/common.sh" ]]; then
      SCRIPT_ROOT="${path}"
      break
    fi
  done

  # HACK(zbehan): We have to fake GCLIENT_ROOT in case we're running inside
  # au_zip enviroment. GCLIENT_ROOT detection became fatal...
  [[ "${SCRIPT_ROOT}" == "${thisdir}" ]] && export GCLIENT_ROOT="."
}

find_common_sh
. "${SCRIPT_ROOT}/common.sh" || exit 1

# Load functions and constants for chromeos-common.
. "${SRC_ROOT}/platform2/installer/share/chromeos-common.sh" &> /dev/null || \
. "/usr/share/misc/chromeos-common.sh" &> /dev/null || \
. "./chromeos-common.sh" || \
  die "Unable to load chromeos-common.sh"

DEFINE_string image "" "The image that should be sent to clients."
DEFINE_string src_image "" "Optional: a source image. If specified, this makes\
 a delta update."
DEFINE_string output "" "Output file"
DEFINE_boolean outside_chroot "${FLAGS_FALSE}" "Running outside of chroot."
DEFINE_string private_key "" "Path to private key in .pem format."
DEFINE_string out_metadata_hash_file "" "Path to output metadata hash file."
DEFINE_boolean extract "${FLAGS_FALSE}" "If set, extract old/new kernel/rootfs \
to [old|new]_[kern|root].dat. Useful for debugging (default: false)"
DEFINE_boolean full_kernel "${FLAGS_FALSE}" "Generate a full kernel update \
even if generating a delta update (default: false)"
DEFINE_integer chunk_size -1 \
  "Delta payload chunk size (default -1 means whole files)"

DEFINE_string src_channel "" "Channel of the src image."
DEFINE_string src_board "" "Board of the src image."
DEFINE_string src_version "" "Version of the src image."
DEFINE_string src_key "" "Key of the src image."
DEFINE_string src_build_channel "" "Channel of the build of the src image."
DEFINE_string src_build_version "" "Version of the build of the src image."

DEFINE_string channel "" "Channel of the target image."
DEFINE_string board "" "Board of the target image."
DEFINE_string version "" "Version of the target image."
DEFINE_string key "" "Key of the target image."
DEFINE_string build_channel "" "Channel of the build of the target image."
DEFINE_string build_version "" "Version of the build of the target image."

# Because we archive/call old versions of this script, we can't easily remove
# command line options, even if we ignore this one now.
DEFINE_boolean patch_kernel "${FLAGS_FALSE}" "Ignored. Present for \
compatibility."

# Specifying any of the following will cause it to not be cleaned up upon exit.
DEFINE_string kern_path "" "File path for extracting the kernel partition."
DEFINE_string root_path "" "File path for extracting the rootfs partition."
DEFINE_string src_kern_path "" \
  "File path for extracting the source kernel partition."
DEFINE_string src_root_path "" \
  "File path for extracting the source rootfs partition."

DEFINE_string work_dir "" "Where to dump temporary files."


# Parse command line
FLAGS "$@" || exit 1
eval set -- "${FLAGS_ARGV}"

# We use this directory to create mount points as a way around crbug.com/358933.
# Use FLAGS_work_dir instead when that workaround is no longer needed.
# When that bug affects us, we will leak temp directories. They should
# be cleaned up on the build servers at the end of cbuildbot.
MOUNTS_ROOT="/tmp/crbug_358933"

SRC_MNT=""
DST_MNT=""
SRC_KERNEL=""
SRC_ROOT=""
DST_KERNEL=""
DST_ROOT=""
STATE_MNT=""

# umount a loopback device, and remove the mount point.
umount_and_rmdir() {
  local mnt_point="$1"

  local umount_ret
  local rmdir_ret=0

  if [[ ! -d  "${mnt_point}" ]]; then
    return 0
  fi

  sudo umount -v "${mnt_point}" || return 1
  sudo rmdir -v "${mnt_point}" || rmdir_ret=$?
  if [[ ${rmdir_ret} -ne 0 ]]; then
    warn "sudo rmdir -v ${mnt_point} exited with: ${rmdir_ret}"
    warn "Please note this failure in crbug.com/358933"
  fi

  # If the rmdir fails, we still return success and keep going. This punts
  # failures up the chain of cleanup code, and makes them responsible for
  # crbug.com/358933.

  return 0
}

cleanup() {
  local err=""

  umount_and_rmdir "${SRC_MNT}" || err=1
  umount_and_rmdir "${DST_MNT}" || err=1
  umount_and_rmdir "${STATE_MNT}" || err=1

  SRC_MNT=""
  DST_MNT=""
  STATE_MNT=""

  if [[ -z "${FLAGS_src_kern_path}" ]]; then
    rm -f "${SRC_KERNEL}" || err=1
  fi
  if [[ -z "${FLAGS_src_root_path}" ]]; then
    rm -f "${SRC_ROOT}" || err=1
  fi
  if [[ -z "${FLAGS_kern_path}" ]]; then
    rm -f "${DST_KERNEL}" || err=1
  fi
  if [[ -z "${FLAGS_root_path}" ]]; then
    rm -f "${DST_ROOT}" || err=1
  fi

  # If we are cleaning up after an error, or if we got an error during
  # cleanup (even if we eventually succeeded) return a non-zero exit
  # code. This triggers additional logging in most environments that call
  # this script.
  if [[ -n "${err}" ]]; then
    die "Cleanup encountered an error."
  fi
}

cleanup_on_error() {
  trap - INT TERM ERR EXIT
  cleanup
  die "Cleanup success after an error."
}

cleanup_on_exit() {
  trap - INT TERM ERR EXIT
  cleanup
}

trap cleanup_on_error INT TERM ERR
trap cleanup_on_exit EXIT

extract_partition_to_temp_file() {
  local filename="$1"
  local partition="$2"
  local temp_file="$3"
  if [[ -z "${temp_file}" ]]; then
    temp_file=$(mktemp --tmpdir="${FLAGS_work_dir}" \
                cros_generate_update_payload.XXXXXX)
  fi
  echo "${temp_file}"

  # Keep `local` decl split from assignment so return code is checked.
  local offset length
  offset=$(partoffset "${filename}" ${partition})  # 512-byte sectors
  length=$(partsize "${filename}" ${partition})  # 512-byte sectors
  local bs=512
  local sectors_per_two_mib=$((2 * 1024 * 1024 / 512))
  if [[ $(( offset % sectors_per_two_mib )) -eq 0 && \
        $(( length % sectors_per_two_mib )) -eq 0 ]]; then
    bs=$((2 * 1024 * 1024))
    offset=$((offset / sectors_per_two_mib))
    length=$((length / sectors_per_two_mib))
  else
    warn "partition offset or length not at 2MiB boundary"
  fi
  dd if="${filename}" of="${temp_file}" bs=${bs} count="${length}" \
      skip="${offset}" 2>/dev/null
}

patch_kernel() {
  local IMAGE="$1"
  local KERN_FILE="$2"

  # Keep `local` decl split from assignment so return code is checked.
  local offset
  offset=$(partoffset "${IMAGE}" 1)
  : $(( offset *= 512 ))
  STATE_MNT=$(mktemp -d --tmpdir="${MOUNTS_ROOT}" state.XXXXXX)
  sudo mount -o ro,loop,offset="${offset}" "${IMAGE}" "${STATE_MNT}"
  dd if="${STATE_MNT}"/vmlinuz_hd.vblock of="${KERN_FILE}" conv=notrunc \
      2>/dev/null
  umount_and_rmdir "${STATE_MNT}"
  STATE_MNT=""
}

extract_kern() {
  local bin_file="$1"
  local kern_out="$2"

  kern_out=$(extract_partition_to_temp_file "${bin_file}" 4 "${kern_out}")
  if $(cmp /dev/zero "${kern_out}" -n 65536 -s); then
    warn "${bin_file}: Kernel B is empty, patching kernel A."
    extract_partition_to_temp_file "${bin_file}" 2 "${kern_out}" > /dev/null
    patch_kernel "${bin_file}" "${kern_out}" >&2
  fi
  echo "${kern_out}"
}

extract_root() {
  local bin_file="$1"
  local root_out="$2"

  extract_partition_to_temp_file "${bin_file}" 3 "${root_out}"
}

mkdir -p "${MOUNTS_ROOT}"

if [[ -n "${FLAGS_src_image}" ]] && \
   [[ "${FLAGS_outside_chroot}" -eq "${FLAGS_FALSE}" ]]; then
  # We need to be in the chroot for generating delta images.
  # by specifying --outside_chroot you can choose not to assert
  # this will allow us to run this script outside chroot.
  # Running this script outside chroot requires copying delta_generator binary
  # and also copying few shared libraries with it.
  assert_inside_chroot
fi

locate_gpt

if [[ "${FLAGS_extract}" -eq "${FLAGS_TRUE}" ]]; then
  if [[ -n "${FLAGS_src_image}" ]]; then
    SRC_KERN_PATH="${FLAGS_src_kern_path:-old_kern.dat}"
    SRC_ROOT_PATH="${FLAGS_src_root_path:-old_root.dat}"
    extract_kern "${FLAGS_src_image}" "${SRC_KERN_PATH}"
    extract_root "${FLAGS_src_image}" "${SRC_ROOT_PATH}"
  fi
  if [[ -n "${FLAGS_image}" ]]; then
    KERN_PATH="${FLAGS_kern_path:-new_kern.dat}"
    ROOT_PATH="${FLAGS_root_path:-new_root.dat}"
    extract_kern "${FLAGS_image}" "${KERN_PATH}"
    extract_root "${FLAGS_image}" "${ROOT_PATH}"
  fi
  echo Done extracting kernel/root
  exit 0
fi

[[ -n "${FLAGS_output}" ]] ||
  die "Error: you must specify an output filename with --output FILENAME"

DELTA="${FLAGS_TRUE}"
PAYLOAD_TYPE="delta"
if [[ -z "${FLAGS_src_image}" ]]; then
  DELTA="${FLAGS_FALSE}"
  PAYLOAD_TYPE="full"
fi

echo "Generating ${PAYLOAD_TYPE} update"

# Sanity check that the real generator exists:
GENERATOR="$(which delta_generator)"
[[ -x "${GENERATOR}" ]] || die "can't find delta_generator"

if [[ "${DELTA}" -eq "${FLAGS_TRUE}" ]]; then
  if [[ "${FLAGS_full_kernel}" -eq "${FLAGS_FALSE}" ]]; then
    SRC_KERNEL=$(extract_kern "${FLAGS_src_image}" "${FLAGS_src_kern_path}")
    echo md5sum of src kernel:
    md5sum "${SRC_KERNEL}"
  else
    echo "Generating a full kernel update."
  fi
  SRC_ROOT=$(extract_root "${FLAGS_src_image}" "${FLAGS_src_root_path}")
  echo md5sum of src root:
  md5sum "${SRC_ROOT}"
fi

DST_KERNEL=$(extract_kern "${FLAGS_image}" "${FLAGS_kern_path}")
DST_ROOT=$(extract_root "${FLAGS_image}" "${FLAGS_root_path}")

DST_MNT=$(mktemp -d --tmpdir="${MOUNTS_ROOT}" dst_root.XXXXXX)
sudo mount -o loop,ro "${DST_ROOT}" "${DST_MNT}"

source "${DST_MNT}"/usr/sbin/write_gpt.sh || warn "Could not read write_gpt.sh"

PARTITION_SIZE_PARAM=
if [[ -n "${ROOTFS_PARTITION_SIZE}" ]]; then
  echo "Using rootfs partition size: ${ROOTFS_PARTITION_SIZE}";
  PARTITION_SIZE_PARAM="-rootfs_partition_size=${ROOTFS_PARTITION_SIZE}"
else
  echo "Using the default partition size."
fi

if [[ "${DELTA}" -eq "${FLAGS_TRUE}" ]]; then
  SRC_MNT=$(mktemp -d --tmpdir="${MOUNTS_ROOT}" src_root.XXXXXX)
  sudo mount -o loop,ro "${SRC_ROOT}" "${SRC_MNT}"

  # Preserve the path during sudo so that unpacked binaries are used when
  # outside the chroot.
  sudo PATH=${PATH} "${GENERATOR}" \
      -new_dir="${DST_MNT}" -new_image="${DST_ROOT}" \
      -new_kernel="${DST_KERNEL}" \
      -old_dir="${SRC_MNT}" -old_image="${SRC_ROOT}" \
      -old_kernel="${SRC_KERNEL}" \
      -out_file="${FLAGS_output}" -private_key="${FLAGS_private_key}" \
      -chunk_size="${FLAGS_chunk_size}" \
      "${PARTITION_SIZE_PARAM}" \
      -old_channel="${FLAGS_src_channel}" -old_board="${FLAGS_src_board}" \
      -old_version="${FLAGS_src_version}" -old_key="${FLAGS_src_key}" \
      -old_build_channel="${FLAGS_src_build_channel}" \
      -old_build_version="${FLAGS_src_build_version}" \
      -new_channel="${FLAGS_channel}" -new_board="${FLAGS_board}" \
      -new_version="${FLAGS_version}" -new_key="${FLAGS_key}" \
      -new_build_channel="${FLAGS_build_channel}" \
      -new_build_version="${FLAGS_build_version}"
else
  "${GENERATOR}" \
      -new_image="${DST_ROOT}" -new_kernel="${DST_KERNEL}" \
      -out_file="${FLAGS_output}" -private_key="${FLAGS_private_key}" \
      "${PARTITION_SIZE_PARAM}" \
      -new_channel="${FLAGS_channel}" -new_board="${FLAGS_board}" \
      -new_version="${FLAGS_version}" -new_key="${FLAGS_key}" \
      -new_build_channel="${FLAGS_build_channel}" \
      -new_build_version="${FLAGS_build_version}"
fi

if [[ -n "${FLAGS_out_metadata_hash_file}" ]]; then
    # The manifest - unfortunately - contain two fields called
    # signature_offset and signature_size with data about the how the
    # manifest is signed. This means we have to pass the signature
    # size used. The value 256 is the number of bytes the SHA-256 hash
    # value of the manifest signed with a 2048-bit RSA key occupies.
    "${GENERATOR}" \
        -in_file="${FLAGS_output}" \
        -signature_size=256 \
        -out_metadata_hash_file="${FLAGS_out_metadata_hash_file}"
fi

echo "Done generating ${PAYLOAD_TYPE} update."