diff --git a/build/debian/build.sh b/build/debian/build.sh
index 7231a7c..9eb478b 100755
--- a/build/debian/build.sh
+++ b/build/debian/build.sh
@@ -10,15 +10,14 @@
 	echo "Builds a debian image and save it to FILE. [sudo is required]"
 	echo "Options:"
 	echo "-h         Print usage and this help message and exit."
-	echo "-a ARCH    Architecture of the image [default is aarch64]"
+	echo "-a ARCH    Architecture of the image [default is host arch: $(uname -m)]"
 	echo "-r         Release mode build"
-	echo "-w         Save temp work directory (for debugging)"
+	echo "-w         Save temp work directory [for debugging]"
 }
 
 check_sudo() {
 	if [ "$EUID" -ne 0 ]; then
-		echo "Please run as root."
-		exit
+		echo "Please run as root." ; exit 1
 	fi
 }
 
@@ -26,17 +25,10 @@
 	while getopts "a:hrw" option; do
 		case ${option} in
 			h)
-				show_help
-				exit;;
+				show_help ; exit
+				;;
 			a)
-				if [[ "$OPTARG" != "aarch64" && "$OPTARG" != "x86_64" ]]; then
-					echo "Invalid architecture: $OPTARG"
-					exit
-				fi
 				arch="$OPTARG"
-				if [[ "$arch" == "x86_64" ]]; then
-					debian_arch="amd64"
-				fi
 				;;
 			r)
 				mode=release
@@ -45,11 +37,21 @@
 				save_workdir=1
 				;;
 			*)
-				echo "Invalid option: $OPTARG"
-				exit
+				echo "Invalid option: $OPTARG" ; exit 1
 				;;
 		esac
 	done
+	case "$arch" in
+		aarch64)
+			debian_arch="arm64"
+			;;
+		x86_64)
+			debian_arch="amd64"
+			;;
+		*)
+			echo "Invalid architecture: $arch" ; exit 1
+			;;
+	esac
 	if [[ "${*:$OPTIND:1}" ]]; then
 		built_image="${*:$OPTIND:1}"
 	fi
@@ -217,7 +219,7 @@
 }
 
 clean_up() {
-	[ "$save_workdir" -eq 0 ] || rm -rf "${workdir}"
+	[ "$save_workdir" -eq 1 ] || rm -rf "${workdir}"
 }
 
 set -e
@@ -230,8 +232,7 @@
 debian_version=bookworm
 config_space=${debian_cloud_image}/config_space/${debian_version}
 resources_dir=${debian_cloud_image}/src/debian_cloud_images/resources
-arch=aarch64
-debian_arch=arm64
+arch="$(uname -m)"
 mode=debug
 save_workdir=0
 
diff --git a/build/debian/build_in_container.sh b/build/debian/build_in_container.sh
index 7fd4c00..5028b74 100755
--- a/build/debian/build_in_container.sh
+++ b/build/debian/build_in_container.sh
@@ -1,35 +1,56 @@
 #!/bin/bash
 
-if [ -z "$ANDROID_BUILD_TOP" ]; then echo "forgot to source build/envsetup.sh?" && exit 1; fi
+show_help() {
+  echo "Usage: sudo $0 [OPTION]..."
+  echo "Builds a debian image and save it to image.raw."
+  echo "Options:"
+  echo "-h         Print usage and this help message and exit."
+  echo "-a ARCH    Architecture of the image [default is host arch: $(uname -m)]"
+  echo "-r         Release mode build"
+  echo "-s         Leave a shell open [default: only if the build fails]"
+  echo "-w         Save temp work directory in the container [for debugging]"
+}
 
-arch=aarch64
+arch="$(uname -m)"
 release_flag=
 save_workdir_flag=
+shell_condition="||"
 
-while getopts "a:rw" option; do
+while getopts "a:rsw" option; do
   case ${option} in
     a)
-      if [[ "$OPTARG" != "aarch64" && "$OPTARG" != "x86_64" ]]; then
-        echo "Invalid architecture: $OPTARG"
-        exit
-      fi
       arch="$OPTARG"
       ;;
+    h)
+      show_help ; exit
+      ;;
     r)
       release_flag="-r"
       ;;
+    s)
+      shell_condition=";"
+      ;;
     w)
       save_workdir_flag="-w"
       ;;
     *)
-      echo "Invalid option: $OPTARG"
-      exit
+      echo "Invalid option: $OPTARG" ; exit 1
       ;;
   esac
 done
 
+if [[ "$arch" != "aarch64" && "$arch" != "x86_64" ]]; then
+  echo "Invalid architecture: $arch" ; exit 1
+fi
+
+if [ -z "$ANDROID_BUILD_TOP" ] ; then
+  echo '`ANDROID_BUILD_TOP` is undefined.'
+  echo 'Please `lunch` an Android target, or manually set the variable.'
+  exit 1
+fi
+
 docker run --privileged -it -v /dev:/dev \
   -v "$ANDROID_BUILD_TOP/packages/modules/Virtualization:/root/Virtualization" \
   --workdir /root/Virtualization/build/debian \
   ubuntu:22.04 \
-  bash -c "/root/Virtualization/build/debian/build.sh -a $arch $release_flag $save_workdir_flag || bash"
+  bash -c "/root/Virtualization/build/debian/build.sh -a $arch $release_flag $save_workdir_flag $shell_condition bash"
diff --git a/build/debian/kokoro/gcp_ubuntu_docker/aarch64/build.sh b/build/debian/kokoro/gcp_ubuntu_docker/aarch64/build.sh
index 130e691..43f0338 100644
--- a/build/debian/kokoro/gcp_ubuntu_docker/aarch64/build.sh
+++ b/build/debian/kokoro/gcp_ubuntu_docker/aarch64/build.sh
@@ -5,7 +5,7 @@
 cd "${KOKORO_ARTIFACTS_DIR}/git/avf/build/debian/"
 sudo losetup -D
 grep vmx /proc/cpuinfo || true
-sudo ./build.sh -r
+sudo ./build.sh -r -a aarch64
 sudo mv images.tar.gz ${KOKORO_ARTIFACTS_DIR} || true
 mkdir -p ${KOKORO_ARTIFACTS_DIR}/logs
 sudo cp -r /var/log/fai/* ${KOKORO_ARTIFACTS_DIR}/logs || true
diff --git a/build/debian/kokoro/gcp_ubuntu_docker/x86_64/build.sh b/build/debian/kokoro/gcp_ubuntu_docker/x86_64/build.sh
index 50ded7b..22ac595 100644
--- a/build/debian/kokoro/gcp_ubuntu_docker/x86_64/build.sh
+++ b/build/debian/kokoro/gcp_ubuntu_docker/x86_64/build.sh
@@ -5,8 +5,7 @@
 cd "${KOKORO_ARTIFACTS_DIR}/git/avf/build/debian/"
 sudo losetup -D
 grep vmx /proc/cpuinfo || true
-sudo ./build.sh -a x86_64 -r
+sudo ./build.sh -r -a x86_64
 sudo mv images.tar.gz ${KOKORO_ARTIFACTS_DIR} || true
-
 mkdir -p ${KOKORO_ARTIFACTS_DIR}/logs
 sudo cp -r /var/log/fai/* ${KOKORO_ARTIFACTS_DIR}/logs || true