fix measure.sh to fit arm64 filepaths

BUG=b/410625960
Change-Id: I507eb62843a514b1c542c22d49d217a4e9e7f56f
Reviewed-on: https://cos-review.googlesource.com/c/third_party/platform/vboot_reference/+/104382
Reviewed-by: Robert Kolchmeyer <rkolchmeyer@google.com>
Tested-by: Cusky Presubmit Bot <presubmit@cos-infra-prod.iam.gserviceaccount.com>
Reviewed-by: Kevin Berry <kpberry@google.com>
diff --git a/scripts/measured_boot_hashes/measure.sh b/scripts/measured_boot_hashes/measure.sh
index af60ab3..d0a6752 100755
--- a/scripts/measured_boot_hashes/measure.sh
+++ b/scripts/measured_boot_hashes/measure.sh
@@ -12,7 +12,7 @@
 declare P12_FILE
 declare VMLINUZ_A_FILE VMLINUZ_B_FILE
 declare GRUB_CFG_FILE
-declare BOOTX64_EFI_FILE GRUB_EFI_FILE
+declare BOOT_EFI_FILE GRUB_EFI_FILE
 declare SHIM_SIG_FILE GRUB_SIG_FILE
 
 
@@ -58,7 +58,7 @@
     VMLINUZ_A_FILE="$TMP_DIR_NAME/vmlinuz.A"
     VMLINUZ_B_FILE="$TMP_DIR_NAME/vmlinuz.B"
     GRUB_CFG_FILE="$TMP_DIR_NAME/grub.cfg"
-    BOOTX64_EFI_FILE="$TMP_DIR_NAME/bootx64.efi"
+    BOOT_EFI_FILE="$TMP_DIR_NAME/boot.efi"
     GRUB_EFI_FILE="$TMP_DIR_NAME/grub-lakitu.efi"
     SHIM_SIG_FILE="$TMP_DIR_NAME/shim_out.sig"
     GRUB_SIG_FILE="$TMP_DIR_NAME/grub_out.sig"
@@ -91,15 +91,32 @@
 # Copies boot-related files from the extracted partition image using mcopy.
 #
 extract_boot_components() {
+    local arch="$1"
+    local shim_path=""
+
+    if [ "$arch" == "x86_64" ]; then
+        shim_path="::/efi/boot/bootx64.efi"
+    elif [ "$arch" == "aarch64" ]; then
+        shim_path="::/efi/boot/bootaa64.efi"
+    else
+        echo "Error: Unknown arch '$arch'."
+        return 1
+    fi
+
     echo "Copying files from partition image '$P12_FILE'..."
+
     declare -A files_to_copy=(
         ["::/syslinux/vmlinuz.A"]="$VMLINUZ_A_FILE"
-        ["::/syslinux/vmlinuz.B"]="$VMLINUZ_B_FILE"
         ["::/efi/boot/grub.cfg"]="$GRUB_CFG_FILE"
-        ["::/efi/boot/bootx64.efi"]="$BOOTX64_EFI_FILE"
+        ["$shim_path"]="$BOOT_EFI_FILE"
         ["::/efi/boot/grub-lakitu.efi"]="$GRUB_EFI_FILE"
     )
 
+    # vmlinuz.B file exists on amd64 but not on arm64
+    if [ "$arch" == "x86_64" ]; then
+        files_to_copy["::/syslinux/vmlinuz.B"]="$VMLINUZ_B_FILE"
+    fi
+
     for src in "${!files_to_copy[@]}"; do
         local dest="${files_to_copy[$src]}"
         echo "Copying '$src' to '$dest'..."
@@ -208,14 +225,15 @@
 
 main() {
     # 1. Parameter Handling & Initial Checks
-    if [[ "$#" -ne 3 ]]; then
-        echo "Usage: $0 <os_image_path> <output_json_file> <channel>"
-        echo "Example: $0 /path/to/image.bin /path/to/output.json stable"
+    if [[ "$#" -ne 4 ]]; then
+        echo "Usage: $0 <os_image_path> <output_json_file> <channel> <build_architecture>"
+        echo "Example: $0 /path/to/image.bin /path/to/output.json stable x86_64"
         return 1
     fi
     local os_image_path="$1"
     local output_json_file="$2"
     local channel="$3"
+    local arch="$4"
 
     if [[ ! -f "$os_image_path" ]]; then
         echo "Error: OS image path '$os_image_path' not found."
@@ -227,24 +245,26 @@
     # 2. Setup and Extraction
     setup_temp_dir
     extract_partition_12 "$os_image_path" || return 1
-    extract_boot_components || return 1
+    extract_boot_components "$arch" || return 1
     
     # 3. Compute All Hashes
     echo "Computing all required hashes..."
-    local vmlinuz_a_hash vmlinuz_b_hash kernel_cmdline_a_hash kernel_cmdline_b_hash shim_hash grub_hash
+    local vmlinuz_a_hash vmlinuz_b_hash="" kernel_cmdline_a_hash kernel_cmdline_b_hash shim_hash grub_hash
 
     vmlinuz_a_hash=$(compute_file_hash "$VMLINUZ_A_FILE")
-    vmlinuz_b_hash=$(compute_file_hash "$VMLINUZ_B_FILE")
+    if [ "$arch" == "x86_64" ]; then
+        vmlinuz_b_hash=$(compute_file_hash "$VMLINUZ_B_FILE")
+    fi
     kernel_cmdline_a_hash=$(compute_cmdline_hash "$GRUB_CFG_FILE" "A")
     kernel_cmdline_b_hash=$(compute_cmdline_hash "$GRUB_CFG_FILE" "B")
-    shim_hash=$(compute_efi_hash "$BOOTX64_EFI_FILE" "$SHIM_SIG_FILE") || return 1
+    shim_hash=$(compute_efi_hash "$BOOT_EFI_FILE" "$SHIM_SIG_FILE") || return 1
     grub_hash=$(compute_efi_hash "$GRUB_EFI_FILE" "$GRUB_SIG_FILE") || return 1
 
     echo "Kernel (vmlinuz.A) hash: $vmlinuz_a_hash"
     echo "Kernel (vmlinuz.B) hash: $vmlinuz_b_hash"
     echo "Kernel cmdline (image A) hash: $kernel_cmdline_a_hash"
     echo "Kernel cmdline (image B) hash: $kernel_cmdline_b_hash"
-    echo "bootx64.efi (shim) hash: $shim_hash"
+    echo "bootx64.efi/bootaa64.efi (shim) hash: $shim_hash"
     echo "grub-lakitu.efi (grub) hash: $grub_hash"
 
     # 4. Final Output