From 9051ebd01996b7dc557ebd071253b4bab48f6872 Mon Sep 17 00:00:00 2001
From: Ibrahim Ansari <ansari.ibrahim1@gmail.com>
Date: Mon, 11 Nov 2024 18:56:37 +0530
Subject: [PATCH] Fix build failure, add disable validation flag

---
 app/dd.go    |  1 +
 app/utils.go | 28 +++++++++++++++++++++++++++-
 main.go      | 29 ++++++++++++++++++-----------
 3 files changed, 46 insertions(+), 12 deletions(-)

diff --git a/app/dd.go b/app/dd.go
index 2f76311..736478b 100644
--- a/app/dd.go
+++ b/app/dd.go
@@ -2,6 +2,7 @@ package app
 
 import (
 	"bufio"
+	"bytes"
 	"errors"
 	"io"
 	"io/fs"
diff --git a/app/utils.go b/app/utils.go
index 59c92ac..a638667 100644
--- a/app/utils.go
+++ b/app/utils.go
@@ -1,6 +1,32 @@
 package app
 
-import "strconv"
+import (
+	"os"
+	"strconv"
+)
+
+type ConfigurationFlags struct {
+	UseSystemDd       bool
+	DisableValidation bool
+}
+
+func ParseCLIFlags() ([]string, ConfigurationFlags) {
+	args := []string{}
+	config := ConfigurationFlags{}
+	if len(os.Args) == 0 {
+		return args, config
+	}
+	for _, arg := range os.Args[1:] {
+		if arg == "--use-system-dd" {
+			config.UseSystemDd = true
+		} else if arg == "--disable-validation" {
+			config.DisableValidation = true
+		} else {
+			args = append(args, arg)
+		}
+	}
+	return args, config
+}
 
 func BytesToString(bytes int, binaryPowers bool) string {
 	i := ""
diff --git a/main.go b/main.go
index 1c3670b..acb7df4 100644
--- a/main.go
+++ b/main.go
@@ -51,25 +51,32 @@ func main() {
 		log.SetFlags(0)
 		log.SetOutput(os.Stderr)
 		log.SetPrefix("[flash] ")
-		if len(os.Args) < 4 {
-			println("Invalid usage: imprint flash <file> <destination> (--use-system-dd)")
+		args, flags := app.ParseCLIFlags()
+		if len(args) < 4 {
+			println("Invalid usage: imprint flash <file> <destination> (--use-system-dd) (--disable-validation)")
 			os.Exit(1)
 		}
-		log.Println("Phase 1/3: Unmounting disk.")
-		if err := app.UnmountDevice(os.Args[3]); err != nil {
+		totalPhases := "3"
+		if flags.DisableValidation {
+			totalPhases = "2"
+		}
+		log.Println("Phase 1/" + totalPhases + ": Unmounting disk.")
+		if err := app.UnmountDevice(args[1]); err != nil {
 			log.Println(err)
-			if !strings.HasSuffix(os.Args[3], "debug.iso") {
+			if !strings.HasSuffix(args[1], "debug.iso") {
 				os.Exit(1)
 			}
 		}
-		log.Println("Phase 2/3: Writing ISO to disk.")
-		if len(os.Args) > 4 && os.Args[4] == "--use-system-dd" {
-			app.RunDd(os.Args[2], os.Args[3])
+		log.Println("Phase 2/" + totalPhases + ": Writing ISO to disk.")
+		if flags.UseSystemDd {
+			app.RunDd(args[0], args[1])
 		} else {
-			app.FlashFileToBlockDevice(os.Args[2], os.Args[3])
+			app.FlashFileToBlockDevice(args[0], args[1])
+		}
+		if flags.DisableValidation {
+			log.Println("Phase 3/" + totalPhases + ": Validating written image on disk.")
+			app.ValidateBlockDeviceContent(args[0], args[1])
 		}
-		log.Println("Phase 3/3: Validating written image on disk.")
-		app.ValidateBlockDeviceContent(os.Args[2], os.Args[3])
 		return
 	}
 	debug := false