#!/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_payload_hash_file "" "Path to output payload hash file."
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_string chunk_size "" \
  "Delta payload chunk size (-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 root_pretruncate_path "" "File path for extracting the rootfs \
partition, pre-truncation."
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}"

STATE_PART_NUM=1
ROOTFS_PART_NUM=3

SRC_KERNEL=""
SRC_ROOT=""
DST_KERNEL=""
DST_ROOT=""

# We include cgpt in the au-generator.zip, so we can call it directly here. We
# don't use chromeos-common.sh versions because we need to not run them as root.

# Usage: gpt_part_offset <image> <partition_number>
# Return the start sector number of the partition number |partition_number| in
# the GPT image |image|.
gpt_part_offset() {
  cgpt show -b -i "$2" "$1"
}

# Usage: gpt_part_size <image> <partition_number>
# Return the size of the partition number |partition_number| in the GPT image
#|image|, in number of sectors.
gpt_part_size() {
  cgpt show -s -i "$2" "$1"
}

cleanup() {
  local err=""

  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=$(gpt_part_offset "${filename}" ${partition})  # 512-byte sectors
  length=$(gpt_part_size "${filename}" ${partition})  # 512-byte sectors
  dd if="${filename}" of="${temp_file}" bs=8M iflag=skip_bytes,count_bytes \
      count="$(( length * 512 ))" skip="$(( offset * 512 ))" status=none
}

patch_kernel() {
  local image="$1"
  local kern_file="$2"
  local err=""

  local state_out vblock
  state_out=$(extract_partition_to_temp_file "${image}" "${STATE_PART_NUM}" "")
  vblock=$(mktemp --tmpdir="${FLAGS_work_dir}" vmlinuz_hd.vblock.XXXXXX)
  e2cp "${state_out}:/vmlinuz_hd.vblock" "${vblock}" || err+="e2cp failed."
  dd if="${vblock}" of="${kern_file}" conv=notrunc status=none ||
    err+="dd failed."
  rm -f "${state_out}" "${vblock}"

  if [[ -n "${err}" ]]; then
    die "Error patching the kernel: ${err}"
  fi
}

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

# ext2fs_size <rootfs>
# Prints the size in bytes of the ext2 filesystem passed in |rootfs|. In case of
# error it returns 1 and doesn't print any output.
ext2fs_size() {
  local rootfs="$1"

  # dumpe2fs is normally installed in /sbin but doesn't require root.
  local PATH="${PATH}:/sbin"
  dumpe2fs "${rootfs}" >/dev/null 2>&1 || return 1
  local fs_blocs fs_blocksize
  fs_blocks=$(dumpe2fs -h "${rootfs}" 2>/dev/null | \
    grep "^Block count:" | cut -f 2 -d :)
  fs_blocksize=$(dumpe2fs -h "${rootfs}" 2>/dev/null | \
    grep "^Block size:" | cut -f 2 -d :)
  echo $(( fs_blocks * fs_blocksize ))
}

# extract_root <bin_file> <root_out> <root_pretruncate_out>
# Extract the rootfs partition from the gpt image |bin_file| and store it in
# |root_out|. If |root_out| is empty, a new temp file will be used. If
# |root_pretruncate_out| is non-empty, saves the pretruncated rootfs partition
# there.  Prints the filename where the rootfs was stored.
extract_root() {
  local bin_file="$1"
  local root_out="$2"
  local root_pretruncate="$3"

  root_out=$(extract_partition_to_temp_file "${bin_file}" "${ROOTFS_PART_NUM}" \
    "${root_out}")

  if [[ -n "${root_pretruncate}" ]]; then
    echo "Saving pre-truncated root as ${root_pretruncate}." >&2
    cp "${root_out}" "${root_pretruncate}"
  fi

  # We only update the filesystem part of the partition, which is stored in the
  # gpt script.
  local root_out_size
  if root_out_size=$(ext2fs_size "${root_out}"); then
    truncate --size="${root_out_size}" "${root_out}"
    echo "Truncated root to ${root_out_size} bytes." >&2
  else
    die "Error truncating the rootfs to the filesystem size."
  fi

  echo "${root_out}"
}

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 [[ "${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}" \
      "${FLAGS_root_pretruncate_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}" \
  "${FLAGS_root_pretruncate_path}")

GENERATOR_ARGS=(
  # Common payload args:
  -major_version=1
  -out_file="${FLAGS_output}"
  -private_key="${FLAGS_private_key}"
  # Target image args:
  -new_image="${DST_ROOT}"
  -new_kernel="${DST_KERNEL}"
  -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}"
)

if [[ "${DELTA}" -eq "${FLAGS_TRUE}" ]]; then
  GENERATOR_ARGS+=(
    # Source image args:
    -old_image="${SRC_ROOT}"
    -old_kernel="${SRC_KERNEL}"
    -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}"
  )

  # The passed chunk_size is only used for delta payload. Use delta_generator's
  # default if no value is provided.
  if [[ -n "${FLAGS_chunk_size}" ]]; then
    echo "Forcing chunk_size to ${FLAGS_chunk_size}";
    GENERATOR_ARGS+=( -chunk_size="${FLAGS_chunk_size}" )
  fi
fi

# Add partition size. Only *required* for minor_version=1.
DST_ROOT_PARTITION_SECTORS=$(gpt_part_size "${FLAGS_image}" \
  "${ROOTFS_PART_NUM}")
if [[ -n "${DST_ROOT_PARTITION_SECTORS}" ]]; then
  DST_ROOT_PARTITION_SIZE=$(( DST_ROOT_PARTITION_SECTORS * 512 ))
  echo "Using rootfs partition size: ${DST_ROOT_PARTITION_SIZE}";
  GENERATOR_ARGS+=( -rootfs_partition_size="${DST_ROOT_PARTITION_SIZE}" )
else
  echo "Using the default partition size."
fi

echo "Running delta_generator with args: ${GENERATOR_ARGS[@]}"
"${GENERATOR}" "${GENERATOR_ARGS[@]}"

if [[ -n "${FLAGS_out_payload_hash_file}" ||
      -n "${FLAGS_out_metadata_hash_file}" ]]; then
    # The out_metadata_hash_file flag requires out_hash_file flag to be set in
    # delta_generator, if the caller doesn't provide it, we set it to /dev/null.
    OUT_PAYLOAD_HASH_FILE="${FLAGS_out_payload_hash_file}"
    if [[ -z "${OUT_PAYLOAD_HASH_FILE}" ]]; then
      OUT_PAYLOAD_HASH_FILE="/dev/null"
    fi
    # 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_hash_file="${OUT_PAYLOAD_HASH_FILE}" \
        -out_metadata_hash_file="${FLAGS_out_metadata_hash_file}"
fi

echo "Done generating ${PAYLOAD_TYPE} update."