From 237aa46dddba2e7875c369b50d754e3267062b46 Mon Sep 17 00:00:00 2001
From: Kyle Isom <kyle@imap.cc>
Date: Thu, 4 May 2023 16:06:56 -0700
Subject: [PATCH] cmd/diskimg: new disk imaging tool

---
 cmd/diskimg/BUILD.bazel |  19 +++++++
 cmd/diskimg/README      |  34 ++++++++++++
 cmd/diskimg/main.go     | 116 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 169 insertions(+)
 create mode 100644 cmd/diskimg/BUILD.bazel
 create mode 100644 cmd/diskimg/README
 create mode 100644 cmd/diskimg/main.go

diff --git a/cmd/diskimg/BUILD.bazel b/cmd/diskimg/BUILD.bazel
new file mode 100644
index 0000000..18a435c
--- /dev/null
+++ b/cmd/diskimg/BUILD.bazel
@@ -0,0 +1,19 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_binary", "go_library")
+
+go_library(
+    name = "diskimg_lib",
+    srcs = ["main.go"],
+    importpath = "git.wntrmute.dev/kyle/goutils/cmd/diskimg",
+    visibility = ["//visibility:private"],
+    deps = [
+        "//ahash",
+        "//dbg",
+        "//die",
+    ],
+)
+
+go_binary(
+    name = "diskimg",
+    embed = [":diskimg_lib"],
+    visibility = ["//visibility:public"],
+)
diff --git a/cmd/diskimg/README b/cmd/diskimg/README
new file mode 100644
index 0000000..7c6d8be
--- /dev/null
+++ b/cmd/diskimg/README
@@ -0,0 +1,34 @@
+diskimg: write disk images
+
+Usage:
+	diskimg [-a algo] [-v] image device
+
+Flags:
+	-a algo		Select the hashing algorithm to use. The default
+			is 'sha256'. Specifying an algorithm of 'list'
+			will print the supported algorithms to standard
+			output and exit with error code 2.
+	-v		Enable verbose (debug) output.
+
+Examples:
+	
+	Copying images/server.img to /dev/sda:
+
+	$ sudo diskimg images/server.img /dev/sda
+
+	Write a bladerunner node image to /dev/sda:
+
+	$ sudo diskimg -v ~/code/bladerunner/packer/build/cm4-cnode-ubuntu-22.04.2.img /dev/sda
+	opening image /home/kyle/code/bladerunner/packer/build/cm4-cnode-ubuntu-22.04.2.img for read
+	/home/kyle/code/bladerunner/packer/build/cm4-cnode-ubuntu-22.04.2.img  416d4c8f890904167419e3d488d097e9c847273376b650546fdb1f6f9809c184
+	opening device /dev/sda for rw
+	writing /home/kyle/code/bladerunner/packer/build/cm4-cnode-ubuntu-22.04.2.img -> /dev/sda
+	wrote 4151312384 bytes to /dev/sda
+	syncing /dev/sda
+	verifying the image was written successfully
+	OK
+
+Motivation:
+
+	I wanted to write something like balena's Etcher, but commandline only.
+
diff --git a/cmd/diskimg/main.go b/cmd/diskimg/main.go
new file mode 100644
index 0000000..f73cb62
--- /dev/null
+++ b/cmd/diskimg/main.go
@@ -0,0 +1,116 @@
+package main
+
+import (
+	"bytes"
+	"flag"
+	"fmt"
+	"io"
+	"os"
+
+	"git.wntrmute.dev/kyle/goutils/ahash"
+	"git.wntrmute.dev/kyle/goutils/dbg"
+	"git.wntrmute.dev/kyle/goutils/die"
+)
+
+const defaultHashAlgorithm = "sha256"
+
+var (
+	hAlgo      string
+	debug = dbg.New()
+)
+
+
+func openImage(imageFile string) (image *os.File, hash []byte, err error) {
+	image, err = os.Open(imageFile)
+	if err != nil {
+		return
+	}
+
+	hash, err = ahash.SumReader(hAlgo, image)
+	if err != nil {
+		return
+	}
+
+	_, err = image.Seek(0, 0)
+	if err != nil {
+		return
+	}
+
+	debug.Printf("%s  %x\n", imageFile, hash)
+	return
+}
+
+func openDevice(devicePath string) (device *os.File, err error) {
+	fi, err := os.Stat(devicePath)
+	if err != nil {
+		return
+	}
+
+	device, err = os.OpenFile(devicePath, os.O_RDWR|os.O_SYNC, fi.Mode())
+	if err != nil {
+		return
+	}
+
+	return
+}
+
+func main() {
+	flag.StringVar(&hAlgo, "a", defaultHashAlgorithm, "default hash algorithm")
+	flag.BoolVar(&debug.Enabled, "v", false, "enable debug logging")
+	flag.Parse()
+
+	if hAlgo == "list" {
+		fmt.Println("Supported hashing algorithms:")
+		for _, algo := range ahash.SecureHashList() {
+			fmt.Printf("\t- %s\n", algo)
+		}
+		os.Exit(2)
+	}
+
+	if flag.NArg() != 2 {
+		die.With("usage: diskimg image device")
+	}
+
+	imageFile := flag.Arg(0)
+	devicePath := flag.Arg(1)
+
+	debug.Printf("opening image %s for read\n", imageFile)
+	image, hash, err := openImage(imageFile)
+	if image != nil {
+		defer image.Close()
+	}
+	die.If(err)
+
+	debug.Printf("opening device %s for rw\n", devicePath)
+	device, err := openDevice(devicePath)
+	if device != nil {
+		defer device.Close()
+	}
+	die.If(err)
+
+	debug.Printf("writing %s -> %s\n", imageFile, devicePath)
+	n, err := io.Copy(device, image)
+	die.If(err)
+	debug.Printf("wrote %d bytes to %s\n", n, devicePath)
+
+	debug.Printf("syncing %s\n", devicePath)
+	err = device.Sync()
+	die.If(err)
+
+	debug.Println("verifying the image was written successfully")
+	_, err = device.Seek(0, 0)
+	die.If(err)
+
+	deviceHash, err := ahash.SumLimitedReader(hAlgo, device, n)
+	die.If(err)
+
+	if !bytes.Equal(deviceHash, hash) {
+		fmt.Fprintln(os.Stderr, "Hash mismatch:")
+		fmt.Fprintf(os.Stderr, "\t%s: %s\n", imageFile, hash)
+		fmt.Fprintf(os.Stderr, "\t%s: %s\n", devicePath, deviceHash)
+		os.Exit(1)
+	}
+
+	debug.Println("OK")
+	os.Exit(0)
+}