From 1a26c6bed3048533d46a99860bb6a44abf89407d Mon Sep 17 00:00:00 2001
From: Florian Loitsch <florian@loitsch.com>
Date: Wed, 6 Nov 2024 17:01:27 +0100
Subject: [PATCH] Add support for extracting an image.

---
 cmd/jag/commands/decode.go         |  2 +-
 cmd/jag/commands/firmware.go       | 57 ++++++++++++++++++++++++++++++
 cmd/jag/commands/firmware_flash.go | 32 +++++++++--------
 3 files changed, 75 insertions(+), 16 deletions(-)

diff --git a/cmd/jag/commands/decode.go b/cmd/jag/commands/decode.go
index 0d25594..26193b7 100644
--- a/cmd/jag/commands/decode.go
+++ b/cmd/jag/commands/decode.go
@@ -216,7 +216,7 @@ func crashDecode(ctx context.Context, envelope string, backtrace string) error {
 		}
 	}
 
-	firmwareElf, err := ExtractFirmware(ctx, sdk, envelopePath, "elf")
+	firmwareElf, err := ExtractFirmware(ctx, sdk, envelopePath, "elf", nil)
 	if err != nil {
 		return err
 	}
diff --git a/cmd/jag/commands/firmware.go b/cmd/jag/commands/firmware.go
index 221e9b8..d0e14f2 100644
--- a/cmd/jag/commands/firmware.go
+++ b/cmd/jag/commands/firmware.go
@@ -43,6 +43,7 @@ func FirmwareCmd() *cobra.Command {
 		},
 	}
 	cmd.AddCommand(FirmwareUpdateCmd())
+	cmd.AddCommand(FirmwareExtractCmd())
 	cmd.Flags().StringP("device", "d", "", "use device with a given name, id, or address")
 	return cmd
 }
@@ -111,3 +112,59 @@ func FirmwareUpdateCmd() *cobra.Command {
 	addFirmwareFlashFlags(cmd, "", "new name of the device, if given")
 	return cmd
 }
+
+func FirmwareExtractCmd() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "extract [envelope]",
+		Short: "Build the firmware image of a Jaguar device",
+		Long: "Build the firmware image of a Jaguar device, for flashing, Wokwi, or QEMU." +
+			"\n" +
+			"Wokwi\n" +
+			"-----\n" +
+			"Remember to use the WiFi SSID \"Wokwi-GUEST\" without a password.\n" +
+			"You need a Club subscription to be able to connect to your device, but\n" +
+			"the device is able to reach the Internet without a subscription.\n" +
+			"\n" +
+			"To run the image go to https://wokwi.com/projects/new/esp32, then\n" +
+			"press F1 and run the command \"Upload Firmware and Start Simulation\".\n",
+		Args:         cobra.MaximumNArgs(1),
+		SilenceUsage: true,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			ctx := cmd.Context()
+
+			output, err := cmd.Flags().GetString("output")
+			if err != nil {
+				return err
+			}
+			if output == "" {
+				return fmt.Errorf("missing output file")
+			}
+
+			return withFirmware(cmd, args, nil, func(newID string, envelopeFile *os.File, config map[string]interface{}) error {
+
+				sdk, err := GetSDK(ctx)
+				if err != nil {
+					return err
+				}
+
+				imageFile, err := ExtractFirmware(ctx, sdk, envelopeFile.Name(), "image", config)
+				if err != nil {
+					return err
+				}
+				defer os.Remove(imageFile.Name())
+
+				imageBin, err := os.ReadFile(imageFile.Name())
+				if err != nil {
+					return err
+				}
+
+				return os.WriteFile(output, imageBin, 0644)
+			})
+		},
+	}
+
+	cmd.Flags().StringP("device", "d", "", "use device with a given name, id, or address")
+	addFirmwareFlashFlags(cmd, "esp32", "name of the device")
+	cmd.Flags().StringP("output", "o", "", "write the firmware image to a file")
+	return cmd
+}
diff --git a/cmd/jag/commands/firmware_flash.go b/cmd/jag/commands/firmware_flash.go
index c00708d..fb62c2b 100644
--- a/cmd/jag/commands/firmware_flash.go
+++ b/cmd/jag/commands/firmware_flash.go
@@ -246,12 +246,12 @@ func ExtractFirmwareBin(ctx context.Context, sdk *SDK, envelopePath string, conf
 	return binaryFile, nil
 }
 
-func ExtractFirmware(ctx context.Context, sdk *SDK, envelopePath string, format string) (*os.File, error) {
+func ExtractFirmware(ctx context.Context, sdk *SDK, envelopePath string, format string, config map[string]interface{}) (*os.File, error) {
 	outputFile, err := os.CreateTemp("", "firmware-"+format+".*")
 	if err != nil {
 		return nil, err
 	}
-	if err := runFirmwareTool(ctx, sdk, envelopePath, "extract", "--format", format, "-o", outputFile.Name()); err != nil {
+	if err := runFirmwareToolWithConfig(ctx, sdk, envelopePath, config, "extract", "--format", format, "-o", outputFile.Name()); err != nil {
 		outputFile.Close()
 		return nil, err
 	}
@@ -263,22 +263,24 @@ func setFirmwareProperty(ctx context.Context, sdk *SDK, envelope *os.File, key s
 }
 
 func runFirmwareToolWithConfig(ctx context.Context, sdk *SDK, envelopePath string, config map[string]interface{}, args ...string) error {
-	configFile, err := os.CreateTemp("", "*.json.config")
-	if err != nil {
-		return err
-	}
-	defer os.Remove(configFile.Name())
+	if config != nil {
+		configFile, err := os.CreateTemp("", "*.json.config")
+		if err != nil {
+			return err
+		}
+		defer os.Remove(configFile.Name())
 
-	configBytes, err := json.Marshal(config)
-	if err != nil {
-		return err
-	}
+		configBytes, err := json.Marshal(config)
+		if err != nil {
+			return err
+		}
 
-	if err := os.WriteFile(configFile.Name(), configBytes, 0666); err != nil {
-		return err
-	}
+		if err := os.WriteFile(configFile.Name(), configBytes, 0666); err != nil {
+			return err
+		}
 
-	args = append(args, "--config", configFile.Name())
+		args = append(args, "--config", configFile.Name())
+	}
 	return runFirmwareTool(ctx, sdk, envelopePath, args...)
 }