update_kernel: Add option to update kernel in the alternate boot slot

The kernel in the non-booting slot can be updated by using the --partition and
--rootfs options. However, this is a little tedious to discover what the
current kernel is and set the options appropriately. Add an option to do
this automatically, for lazy people like me.

BUG=None
TEST=Use --ab_update in combination with the "Chrome OS failsafe kernel
     flashing" doc to boot 5.10 on grunt.

Change-Id: I64da029bac56b9039c2d1c9f5f033209eb1e97a7
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/crosutils/+/2616592
Commit-Queue: Anand K Mistry <amistry@chromium.org>
Tested-by: Anand K Mistry <amistry@chromium.org>
Reviewed-by: Tomasz Figa <tfiga@chromium.org>
Reviewed-by: Nicolas Boichat <drinkcat@chromium.org>
diff --git a/update_kernel.sh b/update_kernel.sh
index 03243a9..54f9230 100755
--- a/update_kernel.sh
+++ b/update_kernel.sh
@@ -29,6 +29,8 @@
 DEFINE_boolean bootonce $FLAGS_FALSE "Mark kernel partition as boot once"
 DEFINE_boolean remote_bootargs $FLAGS_FALSE "Use bootargs from running kernel on target"
 DEFINE_boolean firmware $FLAGS_FALSE "Also update firmwares (/lib/firmware)"
+DEFINE_boolean ab_update $FLAGS_FALSE "Update the kernel in the non-booting \
+kernel slot, similar to an AB update"
 DEFINE_string boot_command "" "Command to run on remote after update (after reboot if applicable)"
 
 ORIG_ARGS=("$@")
@@ -68,6 +70,7 @@
   PARTITION_NUM_KERN_A=2
   PARTITION_NUM_ROOT_A=3
   PARTITION_NUM_KERN_B=4
+  PARTITION_NUM_ROOT_B=5
   PARTITION_NUM_EFI_SYSTEM=12
 }
 
@@ -89,6 +92,15 @@
     REMOTE_VERITY=${FLAGS_FALSE}
     info "System is not using verity: updating firmware and modules"
   fi
+  if [[ ${FLAGS_ab_update} -eq ${FLAGS_TRUE} ]]; then
+    if [[ "${REMOTE_OUT}" == "${FLAGS_device}${PARTITION_NUM_ROOT_A}" ]]; then
+      FLAGS_partition="${FLAGS_device}${PARTITION_NUM_KERN_B}"
+      FLAGS_rootfs="${FLAGS_device}${PARTITION_NUM_ROOT_B}"
+    else
+      FLAGS_partition="${FLAGS_device}${PARTITION_NUM_KERN_A}"
+      FLAGS_rootfs="${FLAGS_device}${PARTITION_NUM_ROOT_A}"
+    fi
+  fi
   if [[ -z "${FLAGS_rootfs}" ]]; then
     FLAGS_rootfs="${REMOTE_OUT}"
   fi