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

# This script takes a path to a rootfs.ext2 which was generated by
# build_image.sh and generates an image that can be used for auto
# update.

# --- BEGIN FACTORY SCRIPTS BOILERPLATE ---
# This script may be executed in a full CrOS source tree or an extracted factory
# bundle with limited tools, so we must always load scripts from $SCRIPT_ROOT
# and search for binary programs in $SCRIPT_ROOT/../bin

SCRIPT="$(readlink -f "$0")"
SCRIPT_ROOT="$(dirname "$SCRIPT")"
. "$SCRIPT_ROOT/lib/cros_image_common.sh" || exit 1
image_find_tool "cgpt" "$SCRIPT_ROOT/../bin"
# --- END FACTORY SCRIPTS BOILERPLATE ---

set -e
# We need 2-3 non-zero parameters.
if [ "$#" -lt 2 ] || [ "$#" -gt 3 ] || [ -z "$1" ] || [ -z "$2" ]; then
  echo "
Usage: $0 kernel_partition_img[:index] rootfs_partition_img[:index] [output_dir]

          Input parameters may be either a simple partition image file, or a
          disk image file name followed by ':' and target partition index number

          If output_dir is omitted, the folder of kernel_partition_img will be
          use.

Examples:
       $0 part_2 part_3
       $0 chromiumos_image.bin:2 part3
       $0 chromiumos_image.bin:2 otherimage.bin:3 /tmp/myoutput
  "
  exit 1
fi

if ! image_has_command pigz; then
  (echo "WARNING:"
   echo " Your system does not have pigz (parallel gzip) installed."
   echo " COMPRESSING WILL BE VERY SLOW. It is recommended to install pigz"
   if image_has_command apt-get; then
     echo " by 'sudo apt-get install pigz'."
   elif image_has_command emerge; then
     echo  " by 'sudo emerge pigz'."
   fi) >&2
fi

if [ $(whoami) = "root" ]; then
  echo "running $0 as root which is unneccessary"
fi

# Usage: load_partition_file VARIABLE_NAME_PREFIX partition_file
# Writes VARIABLE_NAME_PREFIX[, _OFFSE, _SIZE, _SECTORS] by parsing
# partition_file, which can be a simple file or image:partno.
load_partition_file() {
  local var_prefix="$1"
  local var="$2"
  local var_offset=""
  local var_size=""
  local var_sectors=""
  local part_no="${var##*:}"

  # test if var is in image:partno format.
  if [ "$part_no" != "$var" ]; then
    var="${var%:*}"
  else
    part_no=""
  fi

  if [ -z "$part_no" ]; then
    var_offset=0
    var_size="$(stat -c"%s" "$var")" ||
      image_die "Invalid file: $var"
    var_sectors="$((var_size / 512))"
  else
    var_offset="$(image_part_offset "$var" "$part_no")" ||
      image_die "Cannot retieve offset for partition $var:$part_no"
    var_sectors="$(image_part_size "$var" "$part_no")" ||
      image_die "Cannot retieve size for partition $var:$part_no"
    var_size=$((var_sectors * 512))
  fi

  # publish the values
  eval "${var_prefix}"="$var"
  eval "${var_prefix}_OFFSET"="$var_offset"
  eval "${var_prefix}_SIZE"="$var_size"
  eval "${var_prefix}_SECTORS"="$var_sectors"
}

load_partition_file KPART "$1"
load_partition_file ROOTFS "$2"

# Sanity check size.
if [ "$KPART_SIZE" -gt $((16 * 1024 * 1024)) ]; then
  echo "Kernel partition size ($KPART_SIZE bytes) greater than 16 MiB."
  echo "That's too big."
  exit 1
fi

if [ "$#" = "3" ]; then
  FINAL_OUT_DIR="$(readlink -f $3)"
else
  FINAL_OUT_DIR="$(dirname "$(readlink -f $1)")"
fi
FINAL_OUT_FILE="$FINAL_OUT_DIR/update.gz"
echo "Output: $FINAL_OUT_FILE"

# Update payload format:
#  [kernel_size: big-endian uint64][kernel_blob][rootfs_blob]

# Prepare kernel_size by using printf as a number like 00000000003d0900, then
# sed to convert as: \x00\x00\x00\x00\x00\x3d\x09\x00, finally echo -e to
# convert into binary.
KPART_SIZE_SIGNATURE="$(printf "%016x" "$KPART_SIZE" |
                        sed 's/\([0-9a-f][0-9a-f]\)/\\x\1/g')"

# Build the blob!
CS_AND_RET_CODES="$(
  (echo -en "$KPART_SIZE_SIGNATURE"
    echo "Compressing kernel..." >&2
    image_dump_partial_file "$KPART" "$KPART_OFFSET" "$KPART_SECTORS"
    echo "Compressing rootfs..." >&2
    image_dump_partial_file "$ROOTFS" "$ROOTFS_OFFSET" "$ROOTFS_SECTORS") |
  image_gzip_compress -c -9 |
  tee "$FINAL_OUT_FILE" |
  openssl sha1 -binary |
  openssl base64 |
  tr '\n' ' '
  echo ${PIPESTATUS[*]})"

EXPECTED_RET_CODES="0 0 0 0 0 0"
set -- $CS_AND_RET_CODES
CALC_CS="$1"
shift
RET_CODES="$@"
if [ "$RET_CODES" != "$EXPECTED_RET_CODES" ]; then
  echo compression/hash failed. $RET_CODES
  exit 1
fi

echo Success. hash is "$CALC_CS"