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

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_DEV=""
DST_DEV=""
SRC_MNT=""
DST_MNT=""
STATE_MNT=""

PATCHED=$(mktemp -d --tmpdir="${FLAGS_work_dir}" patched.XXXXXX)

# 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 [[ -n "${SRC_DEV}" ]]; then
    loopback_detach "${SRC_DEV}"
  fi

  if [[ -n "${DST_DEV}" ]]; then
    loopback_detach "${DST_DEV}"
  fi

  sudo rm -rf "${PATCHED}" || err=1

  # 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

patch_kernel() {
  local stateful="$1"
  local kern="$2"

  STATE_MNT=$(mktemp -d --tmpdir="${MOUNTS_ROOT}" state.XXXXXX)
  sudo mount -o ro "${stateful}" "${STATE_MNT}"
  dd if="${STATE_MNT}"/vmlinuz_hd.vblock of="${kern}" conv=notrunc 2>/dev/null
  umount_and_rmdir "${STATE_MNT}"
}

extract_kern() {
  local dev="$1"
  local kern_out="$2"
  local kern_a="${dev}p2"
  local kern_b="${dev}p4"
  local stateful="${dev}p1"

  if ! cmp /dev/zero "${kern_b}" -n 65536 -s; then
    if [[ -n "${kern_out}" ]]; then
      cp "${kern_b}" "${kern_out}"
    else
      kern_out="${kern_b}"
    fi
  else
    warn "${bin_file}: Kernel B is empty, patching kernel A."
    [[ -n "${kern_out}" ]] || \
      kern_out=$(mktemp --tmpdir="${PATCHED}" \
                 cros_generate_update_payload.XXXXXX)
    cp "${kern_b}" "${kern_out}"
    patch_kernel "${stateful}" "${kern_out}" >&2
  fi
  echo "${kern_out}"
}

extract_root() {
  local dev="$1"
  local root_out="$2"
  local root_a="${dev}p3"

  if [[ -n "${root_out}" ]]; then
    cp "${root_a}" "${root_out}"
  else
    root_out="${root_a}"
  fi
  echo "${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

if [[ -n "${FLAGS_src_image}" ]]; then
  SRC_DEV=$(loopback_partscan "${FLAGS_src_image}")
  sudo chmod a+r "${SRC_DEV}" "${SRC_DEV}"p*
fi

if [[ -n "${FLAGS_image}" ]]; then
  DST_DEV=$(loopback_partscan "${FLAGS_image}")
  sudo chmod a+r "${DST_DEV}" "${DST_DEV}"p*
fi

if [[ "${FLAGS_extract}" -eq "${FLAGS_TRUE}" ]]; then
  if [[ -n "${FLAGS_src_image}" ]]; then
    extract_kern "${SRC_DEV}" "${FLAGS_src_kern_path:-old_kern.dat}"
    extract_root "${SRC_DEV}" "${FLAGS_src_root_path:-old_root.dat}"
  fi
  if [[ -n "${FLAGS_image}" ]]; then
    extract_kern "${DST_DEV}" "${FLAGS_kern_path:-new_kern.dat}"
    extract_root "${DST_DEV}" "${FLAGS_root_path:-new_root.dat}"
  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 "${SRC_DEV}" "${FLAGS_src_kern_path}")
    echo md5sum of src kernel:
    md5sum "${SRC_KERNEL}"
  else
    echo "Generating a full kernel update."
  fi
  SRC_ROOT=$(extract_root "${SRC_DEV}" "${FLAGS_src_root_path}")
  echo md5sum of src root:
  md5sum "${SRC_ROOT}"
fi

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

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

PARTITION_SIZE_PARAM="-rootfs_partition_size=$(bd_safe_size "${DST_ROOT}")"

if [[ "${DELTA}" -eq "${FLAGS_TRUE}" ]]; then
  SRC_MNT=$(mktemp -d --tmpdir="${MOUNTS_ROOT}" src_root.XXXXXX)
  sudo mount -o 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."