Use common partition handling code

This updates the script to use the common update_partition_table code,
shared with mod_image_for_recovery.

It does not, however, fix the missing support for changing the size of
the stateful partition. Once that is fixed, this script will
automatically start working again.

BUG=chromium-os:35003
TEST=script attempts to build images again
CQ-DEPEND=Ibcd5289389dcadf58ccf0678ecfb29095848b247

Change-Id: I4d2c328f77c70503b996b3faa79decc8f08b1bae
Signed-off-by: Kees Cook <keescook@chromium.org>
Reviewed-on: https://gerrit.chromium.org/gerrit/34681
Reviewed-by: Liam McLoughlin <lmcloughlin@chromium.org>
diff --git a/host/tests/mod_recovery_for_decryption.sh b/host/tests/mod_recovery_for_decryption.sh
index 836d83f..6aa695c 100755
--- a/host/tests/mod_recovery_for_decryption.sh
+++ b/host/tests/mod_recovery_for_decryption.sh
@@ -8,10 +8,9 @@
 # --statefulfs_sectors, and adds the decryption flag file. It is designed
 # to be run from outside of the chroot.
 
-SCRIPT_ROOT=/usr/lib/crosutils
-. "${SCRIPT_ROOT}/common.sh" || exit 1
-. /usr/lib/installer/chromeos-common.sh || exit 1
-locate_gpt
+SCRIPT_ROOT=$(readlink -f $(dirname $(readlink -f "$0"))/../../../../scripts)
+. "${SCRIPT_ROOT}/build_library/build_common.sh" || exit 1
+. "${SCRIPT_ROOT}/build_library/disk_layout_util.sh" || exit 1
 
 assert_inside_chroot
 
@@ -38,50 +37,6 @@
   set -x
 fi
 
-#TODO: this function be moved into a common library (perhaps along
-#      with the gpt functions?)
-update_partition_table() {
-  local src_img=$1          # source image
-  local temp_state=$2       # stateful partition image
-  local resized_sectors=$3  # number of sectors in resized stateful partition
-  local temp_img=$4
-
-  local kern_a_offset=$(partoffset ${src_img} 2)
-  local kern_a_count=$(partsize ${src_img} 2)
-  local kern_b_offset=$(partoffset ${src_img} 4)
-  local kern_b_count=$(partsize ${src_img} 4)
-  local rootfs_offset=$(partoffset ${src_img} 3)
-  local rootfs_count=$(partsize ${src_img} 3)
-  local oem_offset=$(partoffset ${src_img} 8)
-  local oem_count=$(partsize ${src_img} 8)
-  local esp_offset=$(partoffset ${src_img} 12)
-  local esp_count=$(partsize ${src_img} 12)
-
-  local temp_pmbr=$(mktemp -t pmbr.XXXXXX)
-  dd if="${src_img}" of="${temp_pmbr}" bs=512 count=1 &>/dev/null
-
-  trap "rm -rf '${temp_pmbr}'" EXIT
-  # Set up a new partition table
-  install_gpt "${temp_img}" "${rootfs_count}" "${resized_sectors}" \
-    "${temp_pmbr}" "${esp_count}" false \
-    $(((rootfs_count * 512)/(1024 * 1024)))
-
-  # Copy into the partition parts of the file
-  dd if="${src_img}" of="${temp_img}" conv=notrunc bs=512 \
-    seek="${START_ROOTFS_A}" skip=${rootfs_offset} count=${rootfs_count}
-  dd if="${temp_state}" of="${temp_img}" conv=notrunc bs=512 \
-    seek="${START_STATEFUL}"
-  # Copy the full kernel (i.e. with vboot sections)
-  dd if="${src_img}" of="${temp_img}" conv=notrunc bs=512 \
-    seek="${START_KERN_A}" skip=${kern_a_offset} count=${kern_a_count}
-  dd if="${src_img}" of="${temp_img}" conv=notrunc bs=512 \
-    seek="${START_KERN_B}" skip=${kern_b_offset} count=${kern_b_count}
-  dd if="${src_img}" of="${temp_img}" conv=notrunc bs=512 \
-    seek="${START_OEM}" skip=${oem_offset} count=${oem_count}
-  dd if="${src_img}" of="${temp_img}" conv=notrunc bs=512 \
-    seek="${START_ESP}" skip=${esp_offset} count=${esp_count}
-}
-
 resize_stateful() {
   # Rebuild the image with larger stateful.
   local err=0