From 89e625791e53feee69144d99c7c1618505087e67 Mon Sep 17 00:00:00 2001
From: Frieder Paape <frieder@konvera.io>
Date: Sat, 16 Nov 2024 06:52:33 +0700
Subject: [PATCH] chore: download image from release if not specified

---
 README.md    |  8 ++++----
 lib/bm.sh    | 30 ++++++++++++++++++++++++++----
 lib/cloud.sh | 43 +++++++++++++++++++++++++++++++++++--------
 3 files changed, 65 insertions(+), 16 deletions(-)
 mode change 100755 => 100644 lib/cloud.sh

diff --git a/README.md b/README.md
index 6562ed7..cf96823 100644
--- a/README.md
+++ b/README.md
@@ -15,16 +15,16 @@ One command, `./zap` - and you've got yourself a TDX box.
 2. Deploy the flashbox VM:
 ```bash
 # Local deployment (non-TDX)
-./zap --mode normal --image flashbox.raw
+./zap --mode normal
 
 # Local deployment (TDX)
-./zap --mode tdx --image flashbox.raw
+./zap --mode tdx
 
 # Azure deployment
-./zap azure myvm eastus flashbox.azure.vhd
+./zap azure myvm eastus
 
 # GCP deployment
-./zap gcp myvm us-east4 flashbox.tar.gz
+./zap gcp myvm us-east4
 ```
 
 ### Known Issues
diff --git a/lib/bm.sh b/lib/bm.sh
index 244a57b..a00a37a 100644
--- a/lib/bm.sh
+++ b/lib/bm.sh
@@ -1,10 +1,12 @@
 #!/bin/bash
 
+# This script can run both standard QEMU VMs and TDX-enabled VMs
+
 usage() {
     echo "Usage: $0 [options]"
     echo "Options:"
     echo "  --mode <normal|tdx>       VM mode (default: normal)"
-    echo "  --image PATH              Path to VM image (required)"
+    echo "  --image PATH              Path to VM image (optional, will download flashbox.raw if not provided)"
     echo "  --ram SIZE                RAM size in GB (default: 32)"
     echo "  --cpus NUMBER             Number of CPUs (default: 16)"
     echo "  --ssh-port PORT           SSH port forwarding (default: 10022)"
@@ -25,6 +27,25 @@ cleanup() {
     sleep 3
 }
 
+download_flashbox() {
+    if [ -f "flashbox.raw" ]; then
+        echo "Using existing flashbox.raw"
+    else
+        echo "Downloading flashbox.raw..."
+        DOWNLOAD_URL=$(curl -s https://api.github.com/repos/flashbots/flashbox/releases/latest | grep "browser_download_url.*flashbox\.raw" | cut -d '"' -f 4)
+        if [ -z "$DOWNLOAD_URL" ]; then
+            echo "Error: Could not find download URL for flashbox.raw"
+            exit 1
+        fi
+        wget "$DOWNLOAD_URL" || {
+            echo "Error: Failed to download flashbox.raw"
+            exit 1
+        }
+    fi
+    echo "flashbox.raw is ready"
+    VM_IMG="flashbox.raw"
+}
+
 # Default values
 MODE="normal"
 RAM_SIZE="32"
@@ -34,6 +55,7 @@ ADDITIONAL_PORTS=""
 PROCESS_NAME="qemu-vm"
 LOGFILE="/tmp/qemu-guest.log"
 OVMF_PATH="/usr/share/ovmf/OVMF.fd"
+VM_IMG=""
 
 # Parse command line arguments
 while [[ $# -gt 0 ]]; do
@@ -84,10 +106,9 @@ while [[ $# -gt 0 ]]; do
     esac
 done
 
-# Check required parameters
+# If no image path provided, download flashbox.raw
 if [ -z "$VM_IMG" ]; then
-    echo "Error: VM image path is required"
-    usage
+    download_flashbox
 fi
 
 # Verify mode
@@ -130,6 +151,7 @@ QEMU_CMD="qemu-system-x86_64 -D $LOGFILE \
     -cpu host \
     -nographic \
     -nodefaults \
+    -daemonize \
     ${PORT_FORWARDS} \
     -drive file=${VM_IMG},if=none,id=virtio-disk0 \
     -device virtio-blk-pci,drive=virtio-disk0 \
diff --git a/lib/cloud.sh b/lib/cloud.sh
old mode 100755
new mode 100644
index 09f2a54..f93e9bc
--- a/lib/cloud.sh
+++ b/lib/cloud.sh
@@ -14,8 +14,8 @@ usage() {
     echo ""
     echo "Arguments:"
     echo "  name                      Resource name/prefix for the deployment"
-    echo "  region                    Cloud region to deploy in (default: eastus for Azure, us-east4 for GCP)"
-    echo "  image-path               Path to VM image (required for deploy)"
+    echo "  region                    Cloud region to deploy in (default: westeurope for Azure, us-east4 for GCP)"
+    echo "  image-path               Path to VM image (optional, will download appropriate image if not provided)"
     echo ""
     echo "Options:"
     echo "  --machine-type TYPE      VM size (default: Standard_EC4eds_v5 for Azure, c3-standard-4 for GCP)"
@@ -24,6 +24,37 @@ usage() {
     exit 1
 }
 
+download_flashbox() {
+    local cloud=$1
+    local image_name
+    local expected_file
+    
+    if [[ "$cloud" == "azure" ]]; then
+        image_name="flashbox.azure.vhd"
+        expected_file="$image_name"
+    else
+        image_name="flashbox.raw.tar.gz"
+        expected_file="$image_name"
+    fi
+
+    if [ -f "$expected_file" ]; then
+        echo "Using existing $expected_file"
+    else
+        echo "Downloading $image_name..."
+        local DOWNLOAD_URL=$(curl -s https://api.github.com/repos/flashbots/flashbox/releases/latest | grep "browser_download_url.*${image_name}" | cut -d '"' -f 4)
+        if [ -z "$DOWNLOAD_URL" ]; then
+            echo "Error: Could not find download URL for $image_name"
+            exit 1
+        fi
+        wget "$DOWNLOAD_URL" || {
+            echo "Error: Failed to download $image_name"
+            exit 1
+        }
+    fi
+    echo "$expected_file is ready"
+    echo "$expected_file"
+}
+
 check_dependencies() {
     local cloud=$1
     if [[ "$cloud" == "azure" ]]; then
@@ -98,9 +129,6 @@ create_azure_deployment() {
     # Create NSG with base rules
     echo "Creating network security group..."
     az network nsg create --name "$name" --resource-group "$name" --location "$region"
-
-    # Add a small delay to ensure NSG is fully created
-    sleep 5
     
     # Add SSH rule with optional IP restriction
     local ssh_source="${ssh_source_ip:-*}"
@@ -277,11 +305,10 @@ fi
 # Execute command
 case $COMMAND in
     deploy)
+        check_dependencies "$CLOUD"
         if [[ -z "$IMAGE_PATH" ]]; then
-            echo "Error: Image path required for deploy command"
-            usage
+            IMAGE_PATH=$(download_flashbox "$CLOUD")
         fi
-        check_dependencies "$CLOUD"
         if [[ "$CLOUD" == "azure" ]]; then
             create_azure_deployment "$NAME" "$REGION" "$IMAGE_PATH" "$MACHINE_TYPE" "$SSH_SOURCE_IP" "$ADDITIONAL_PORTS"
         else