#!/bin/bash

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

# modify_netboot_ip.sh --image=[image] --addr=[ip address]
#
# This script takes initrd image generated by make_netboot.sh as input, and
# modify mini-Omaha server address inside lsb-factory.

. "$(dirname "$(readlink -f "$0")")/factory_common.sh" || exit 1

DEFINE_string image "" "Path to the initrd image"
DEFINE_string addr "" "New mini-Omaha server address"

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

[ -z "${FLAGS_image}" ] && die "No image specified"
[ -z "${FLAGS_addr}" ] && die "No server address specified"

cleanup() {
  info "Cleaning up"
  sudo umount ${WORK_DIR}/rootfs_dir >/dev/null 2>&1 || true
  rm -f "${WORK_DIR}/r.gz" "${WORK_DIR}/r" "${IMAGE_INFO}" >/dev/null 2>&1
  rmdir "${WORK_DIR}/rootfs_dir" "${WORK_DIR}" >/dev/null 2>&1 || true
}

set -e
trap cleanup EXIT

IMAGE_PATH="$(readlink -f ${FLAGS_image})"
WORK_DIR="$(mktemp -d modip_XXXXXXXX | xargs readlink -f)"
IMAGE_INFO="${WORK_DIR}/image_info"

# Save image metadata from header so that we can restore it later.
info "Saving image metadata from ${IMAGE_PATH}"
mkimage -l "${IMAGE_PATH}" >"${IMAGE_INFO}" || die

# Header resides in the first 64 bytes. Let's remove it.
info "Removing UImage header"
dd if="${IMAGE_PATH}" of="${WORK_DIR}/r.gz" bs=64 skip=1 >/dev/null 2>&1

# initrd image generated by make_netboot is gzip'ed. Decompresses it and modify
# lsb-factory inside.
info "Decompressing image"
gunzip ${WORK_DIR}/r.gz
mkdir ${WORK_DIR}/rootfs_dir
info "Modifying server address to ${FLAGS_addr}"
sudo mount -o loop ${WORK_DIR}/r ${WORK_DIR}/rootfs_dir
LSB_FACTORY_PATH="rootfs_dir/mnt/stateful_partition/dev_image/etc/lsb-factory"
sudo sed -i "s,http://[^:]*:,http://${FLAGS_addr}:," \
    "${WORK_DIR}/${LSB_FACTORY_PATH}"
sync
sudo umount ${WORK_DIR}/rootfs_dir
info "Compressing image"
gzip -9 ${WORK_DIR}/r

info "Append UImage header"
IMAGE_NAME=$(grep "Image Name" "${IMAGE_INFO}" | sed 's/^[^:]*: *//')
LOAD_ADDR="0x$(grep "Load Address" "${IMAGE_INFO}" | sed 's/^[^:]*: *//')"
ENTRY_ADDR="0x$(grep "Entry Point" "${IMAGE_INFO}" | sed 's/^[^:]*: *//')"
if grep -q "x86" "${IMAGE_INFO}" ; then
  ARCH="x86"
else
  ARCH="arm"
fi
mkimage -A "${ARCH}" -O linux -T ramdisk -a "${LOAD_ADDR}" -e "${ENTRY_ADDR}" \
    -n "${IMAGE_NAME}" -C gzip -d "${WORK_DIR}/r.gz" "${IMAGE_PATH}" >/dev/null

info "Done"