#!/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-install
# NOTE: Needs to be called from outside the chroot.
. "/usr/lib/installer/chromeos-common.sh" &> /dev/null || \
. "${SRC_ROOT}/platform2/installer/share/chromeos-common.sh" &> /dev/null || \
. "./chromeos-common.sh" || \
  die "Unable to load /usr/lib/installer/chromeos-common.sh"

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

# umount a loopback device, and remove the mount point.
# This can be called normally (set -e), or during cleanup (set +3).
umount_and_rmdir() {
  local mnt_point="$1"

  local umount_ret
  local rmdir_ret=0

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


  for x in {1..30}; do
    umount_ret=0
    sudo umount -v "$mnt_point" || umount_ret=$?
    if [ ${umount_ret} -ne 0 ]; then
      echo "sudo umount -vl ${mnt_point} exited with: $umount_ret"
      echo "Please note this failure in crbug.com/358933"
    fi

    rmdir_ret=0
    sudo rmdir -v "$mnt_point" || rmdir_ret=$?
    if [ ${rmdir_ret} -eq 0 ]; then
      # This means we worked!
      return 0
    else
      echo "sudo rmdir -v ${mnt_point} exited with: $rmdir_ret"
      echo "Please note this failure in crbug.com/358933"
    fi

    sleep 60
    echo "Retrying umount_and_rmdir ${mnt_point}"
  done

  return 1
}

# Pass any argument to this method during normal cleanup.
# Otherwise, it will assume it was called from a trap, and exit 1.
cleanup() {
  local normally_invoked="$1"
  local err=""

  set +e

  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
    echo "Cleanup encountered an error."
    exit 2
  fi

  if [[ -z "$normally_invoked" ]]; then
    echo "Cleanup success after an error."
    exit 1
  fi
}

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 -a \
       $(( $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
}

extract_kern_root() {
  local bin_file="$1"
  local kern_out="$2"
  local root_out="$3"

  if [ -z "$kern_out" ]; then
    die "missing kernel output filename"
  fi
  if [ -z "$root_out" ]; then
    die "missing root output filename"
  fi

  extract_partition_to_temp_file "$bin_file" 4 "$kern_out" > /dev/null
  extract_partition_to_temp_file "$bin_file" 3 "$root_out" > /dev/null
}

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}"

set -e

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_root "$FLAGS_src_image" "$SRC_KERN_PATH" "$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_root "$FLAGS_image" "$KERN_PATH" "$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"

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

  echo "Patching kernel" $KERN_FILE
  echo "   into" $IMAGE
  # Keep `local` decl split from assignment so return code is checked.
  local offset
  offset=$(partoffset "${IMAGE}" 1)
  : $(( offset *= 512 ))
  STATE_MNT=$(mktemp -d --tmpdir="${FLAGS_work_dir}" 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=""
}

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

trap cleanup INT TERM EXIT
if [ "$DELTA" -eq "$FLAGS_TRUE" ]; then
  if [ "$FLAGS_full_kernel" -eq "$FLAGS_FALSE" ]; then
    SRC_KERNEL=$(extract_partition_to_temp_file "$FLAGS_src_image" 4 \
                 "$FLAGS_src_kern_path")
    if $(cmp /dev/zero "${SRC_KERNEL}" -n 65536 -s); then
      echo "Source kernel B is empty, patching kernel A."
      rm -rf "$SRC_KERNEL"
      SRC_KERNEL=$(extract_partition_to_temp_file "$FLAGS_src_image" 2 \
                   "$FLAGS_src_kern_path")
      patch_kernel "$FLAGS_src_image" "$SRC_KERNEL"
    fi
    echo md5sum of src kernel:
    md5sum "$SRC_KERNEL"
  else
    echo "Generating a full kernel update."
  fi
  SRC_ROOT=$(extract_partition_to_temp_file "$FLAGS_src_image" 3 \
             "$FLAGS_src_root_path")

  echo md5sum of src root:
  md5sum "$SRC_ROOT"
fi

DST_KERNEL=$(extract_partition_to_temp_file "$FLAGS_image" 4 \
             "$FLAGS_kern_path")
if $(cmp /dev/zero "${DST_KERNEL}" -n 65536 -s); then
  echo "Destination kernel B is empty, patching kernel A."
  rm -rf "$DST_KERNEL"
  DST_KERNEL=$(extract_partition_to_temp_file "$FLAGS_image" 2 \
               "$FLAGS_kern_path")
  patch_kernel "$FLAGS_image" "$DST_KERNEL"
fi
DST_ROOT=$(extract_partition_to_temp_file "$FLAGS_image" 3 \
           "$FLAGS_root_path")

DST_MNT=$(mktemp -d --tmpdir="${FLAGS_work_dir}" src_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="${FLAGS_work_dir}" 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

trap - INT TERM EXIT
cleanup noexit

echo "Done generating $PAYLOAD_TYPE update."