From 947660b7733a4d062d9fc14d8929ea5216b05510 Mon Sep 17 00:00:00 2001
From: Nils <nils@byreqz.de>
Date: Tue, 8 Feb 2022 12:28:49 +0100
Subject: [PATCH] add output verification

---
 README.md |  1 +
 main.go   | 68 ++++++++++++++++++++++++++++++++++++++-----------------
 2 files changed, 48 insertions(+), 21 deletions(-)

diff --git a/README.md b/README.md
index bbd05c9..652bc1d 100644
--- a/README.md
+++ b/README.md
@@ -12,6 +12,7 @@ arguments:
   -d, --device string   target device
   -f, --force           override safety features
   -i, --input string    input file
+  -n, --no-hash         disable hash verification
 ```
 
 If no image or device is specified, etcher will enter interactive mode and prompt for the missing parameters.
diff --git a/main.go b/main.go
index e1f377c..f8cf300 100644
--- a/main.go
+++ b/main.go
@@ -8,6 +8,7 @@ import (
 	"strings"
 	"runtime"
 	"strconv"
+	"crypto/sha256"
 	"github.com/schollz/progressbar/v3"
 	"github.com/fatih/color"
 	"github.com/briandowns/spinner"
@@ -18,11 +19,13 @@ import (
 var device string
 var input string
 var force bool
+var disable_hash bool
 
 func init() {
 	flag.StringVarP(&device, "device", "d", "", "target device")
 	flag.StringVarP(&input, "input", "i", "", "input file")
 	flag.BoolVarP(&force, "force", "f", false, "override safety features")
+	flag.BoolVarP(&disable_hash, "no-hash", "n", false, "disable hash verification")
 	flag.Parse()
 }
 
@@ -88,26 +91,6 @@ func WriteImage(image *os.File, target *os.File, size int64) (int64, error) {
 	return written, err
 }
 
-func Sync(image *os.File, target *os.File) error {
-	err := image.Sync()
-	if err != nil {
-		return err
-	}
-	err = target.Sync()
-	if err != nil {
-		return err
-	}
-	err = image.Close()
-	if err != nil {
-		return err
-	}
-	err = target.Close()
-	if err != nil {
-		return err
-	}
-	return nil
-}
-
 func PrintAvail() {
 	if runtime.GOOS == "linux" {
 		block, _ := os.ReadDir("/sys/block")
@@ -216,6 +199,17 @@ func main() {
 		targetisblock = true
 		_, _ = target.Seek(0, 0)
 	}
+	prehash := sha256.New()
+	if ! (force || disable_hash) {
+		if err != nil {
+			s.Stop()
+			fmt.Println("\r[", color.RedString("✘"), "]  Getting file details                     ")
+			log.Fatal(err)
+		}
+
+		_, err = io.Copy(prehash, image)
+		_, _ = image.Seek(0, 0)
+	}
 	if err != nil {
 		s.Stop()
 		fmt.Println("\r[", color.RedString("✘"), "]  Getting file details                     ")
@@ -253,6 +247,7 @@ func main() {
 		}
 	}
 	written, err := WriteImage(image, target, inputsize)
+	_, _ = target.Seek(0, 0)
 	if err != nil {
 		fmt.Println("\r[", color.RedString("✘"), "]  Writing image,", written, "bytes written                                                    ")
 		log.Fatal(err)
@@ -263,7 +258,13 @@ func main() {
 	s.Prefix = "[ "
 	s.Suffix = " ]  Syncing"
 	s.Start()
-	err = Sync(image, target)
+	err = image.Sync()
+	if err != nil {
+		s.Stop()
+		fmt.Println("\r[", color.RedString("✘"), "]  Syncing                    ")
+		log.Fatal(err)
+	}
+	err = target.Sync()
 	if err != nil {
 		s.Stop()
 		fmt.Println("\r[", color.RedString("✘"), "]  Syncing                    ")
@@ -272,4 +273,29 @@ func main() {
 		s.Stop()
 		fmt.Println("\r[", color.GreenString("✓"), "]  Syncing                  ")
 	}
+	if ! (force || disable_hash) {
+		s.Prefix = "[ "
+		s.Suffix = " ]  Verifying"
+		s.Start()
+		posthash := sha256.New()
+		_, err = io.CopyN(posthash, target, inputsize)
+		presum := fmt.Sprintf("%x", prehash.Sum(nil))
+		postsum := fmt.Sprintf("%x", posthash.Sum(nil))
+		if err != nil || presum != postsum {
+			s.Stop()
+			fmt.Println("\r[", color.RedString("✘"), "]  Verifying                    ")
+			log.Fatal(err)
+		} else {
+			s.Stop()
+			fmt.Println("\r[", color.GreenString("✓"), "]  Verifying                  ")
+		}
+	}
+	err = image.Close()
+	if err != nil {
+		log.Fatal(err)
+	}
+	err = target.Close()
+	if err != nil {
+		log.Fatal(err)
+	}
 }