From 0625418d9d2cf618e7f689fe907680cd9b718425 Mon Sep 17 00:00:00 2001 From: Nils Date: Sun, 8 Oct 2023 11:57:13 +0200 Subject: [PATCH] restructure packages and start picking apart functions --- cmd/cetcher/main.go | 283 ++++++++++++++++++++++++++++++++++++++ etcher/etcher.go | 11 ++ etcher/etcher_linux.go | 68 ++++++++++ main.go | 301 +---------------------------------------- 4 files changed, 365 insertions(+), 298 deletions(-) create mode 100644 cmd/cetcher/main.go create mode 100644 etcher/etcher.go create mode 100644 etcher/etcher_linux.go diff --git a/cmd/cetcher/main.go b/cmd/cetcher/main.go new file mode 100644 index 0000000..232e332 --- /dev/null +++ b/cmd/cetcher/main.go @@ -0,0 +1,283 @@ +// package cetcher is the cli app for the go-etcher project +package cetcher + +import ( + "crypto/sha256" + "fmt" + "io" + "log" + "os" + "strings" + "time" + + ac "github.com/JoaoDanielRufino/go-input-autocomplete" + "github.com/briandowns/spinner" + "github.com/byReqz/go-etcher/etcher" + "github.com/fatih/color" + "github.com/schollz/progressbar/v3" + flag "github.com/spf13/pflag" +) + +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() +} + +func GetPath() string { + path, err := ac.Read("[ " + color.YellowString("i") + " ] Please input your image file: ") + if err != nil { + log.Fatal(err) + } + path = strings.TrimSpace(path) + if strings.HasPrefix(path, "~/") { + homedir, err := os.UserHomeDir() + if err != nil { + log.Fatal(err) + } + return homedir + path[1:] + } + if path == "" { + fmt.Println("[", color.RedString("!"), "] No image given, retrying.") + path = GetPath() + } + return path +} + +func GetDest() string { + PrintAvail() + dest, err := ac.Read("[ " + color.YellowString("i") + " ] Please input destination: ") + if err != nil { + log.Fatal(err) + } + dest = strings.TrimSpace(dest) + if strings.HasPrefix(dest, "~/") { + homedir, err := os.UserHomeDir() + if err != nil { + log.Fatal(err) + } + return homedir + dest[1:] + } + if dest == "" { + fmt.Println("[", color.RedString("!"), "] No destination given, retrying.") + dest = GetDest() + } + return dest +} + +func WriteImage(image *os.File, target *os.File, size int64) (int64, error) { + bar := progressbar.NewOptions(int(size), + progressbar.OptionSetWriter(os.Stderr), + progressbar.OptionShowBytes(true), + progressbar.OptionSetWidth(25), + progressbar.OptionSetDescription("Writing image file..."), + progressbar.OptionSetTheme(progressbar.Theme{ + Saucer: "=", + SaucerHead: ">", + SaucerPadding: " ", + BarStart: "[", + BarEnd: "]", + })) + writer := io.MultiWriter(target, bar) + written, err := io.Copy(writer, image) + if err != nil { + return 0, err + } + return written, err +} + +func PrintAvail() { + devices, err := etcher.GetBlockdevices() + if err != nil { + log.Fatal(err) + } + + fmt.Println("Available devices:") + + for _, dev := range devices { + fmt.Print(" * ", dev.Path) + if dev.Size > 0 { + fmt.Print(" [", (dev.Size / 1024 / 1024 / 1024), "GB]\n") + } else { + fmt.Println("") + } + } +} + +func Main() { + s := spinner.New(spinner.CharSets[14], 100*time.Millisecond) + + if input == "" { + if len(flag.Args()) == 0 { + input = GetPath() + } else if len(flag.Args()) > 0 { + input = flag.Args()[0] + } + } + + if device == "" { + if len(flag.Args()) == 0 { + device = GetDest() + } else if len(flag.Args()) > 0 { + if input == flag.Args()[0] && len(flag.Args()) > 1 { + device = flag.Args()[1] + } else if input != flag.Args()[0] && len(flag.Args()) > 0 { + device = flag.Args()[0] + } + } + } + + s.Prefix = "[ " + s.Suffix = " ] Getting file details" + s.Start() + statinput, err := os.Stat(input) + if err != nil { + s.Stop() + fmt.Println("\r[", color.RedString("✘"), "] Getting file details ") + log.Fatal(err) + } + statdevice, err := os.Stat(device) + if err != nil { + s.Stop() + fmt.Println("\r[", color.RedString("✘"), "] Getting file details ") + log.Fatal(err) + } + image, err := os.Open(input) + if err != nil { + s.Stop() + fmt.Println("\r[", color.RedString("✘"), "] Getting file details ") + log.Fatal(err) + } + var inputsize int64 + var inputisblock bool + if statinput.Size() != 0 { + inputsize = statinput.Size() + inputisblock = false + } else { + inputsize, _ = image.Seek(0, io.SeekEnd) + inputisblock = true + _, _ = image.Seek(0, 0) + } + target, err := os.OpenFile(device, os.O_RDWR, 0660) + if err != nil { + s.Stop() + fmt.Println("\r[", color.RedString("✘"), "] Getting file details ") + log.Fatal(err) + } + var targetsize int64 + var targetisblock bool + if statdevice.Size() != 0 { + targetsize = statdevice.Size() + targetisblock = false + } else { + targetsize, err = target.Seek(0, io.SeekEnd) + 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 ") + log.Fatal(err) + } else { + s.Stop() + fmt.Println("\r[", color.GreenString("✓"), "] Getting file details ") + } + inputmb := fmt.Sprint("[", inputsize/1024/1024, "MB]") + devicemb := fmt.Sprint("[", targetsize/1024/1024, "MB]") + var inputblock string + var targetblock string + if inputisblock { + inputblock = "[Blockdevice]" + } else { + inputblock = "[File]" + } + if targetisblock { + targetblock = "[Blockdevice]" + } else { + targetblock = "[File]" + } + fmt.Println("[", color.BlueString("i"), "] Input device/file: "+input, inputmb, inputblock) + fmt.Println("[", color.BlueString("i"), "] Output device/file: "+device, devicemb, targetblock) + if !force { + if inputsize > targetsize { + fmt.Println("[", color.RedString("w"), "]", color.RedString(" Warning:"), "Input file seems to be bigger than the destination!") + } + fmt.Print(color.HiWhiteString("Do you want to continue? [y/N]: ")) + var yesno string + _, _ = fmt.Scanln(&yesno) + yesno = strings.TrimSpace(yesno) + if !(yesno == "y" || yesno == "Y") { + log.Fatal("aborted") + } + } + 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) + } else { + fmt.Println("\r[", color.GreenString("✓"), "] Writing image,", written, "bytes written ") + } + + s.Prefix = "[ " + s.Suffix = " ] Syncing" + s.Start() + 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 ") + log.Fatal(err) + } else { + 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) + } +} diff --git a/etcher/etcher.go b/etcher/etcher.go new file mode 100644 index 0000000..0aff848 --- /dev/null +++ b/etcher/etcher.go @@ -0,0 +1,11 @@ +package etcher + +type Blockdevice struct { + Path string + Size int + Type string + //ID string +} + +// CheckHash checks if the given file has been properly applied to the given blockdevice. +//func CheckHash(f os.File, b Blockdevice) error {} diff --git a/etcher/etcher_linux.go b/etcher/etcher_linux.go new file mode 100644 index 0000000..3d34a00 --- /dev/null +++ b/etcher/etcher_linux.go @@ -0,0 +1,68 @@ +package etcher + +import ( + "fmt" + "io" + "os" + "strconv" + "strings" +) + +// getBlockdeviceSize gets the size of a blockdevice from the kernel filesystem. +func getBlockdeviceSize(name string) (size int, err error) { + sizefile, err := os.Open("/sys/block/" + name + "/size") + if err != nil { + return + } + sizeread, err := io.ReadAll(sizefile) + if err != nil { + return + } + _ = sizefile.Close() + + sizestring := strings.TrimSuffix(string(sizeread), "\n") + size, err = strconv.Atoi(sizestring) + if err != nil { + return + } + size = size * 512 // multiply blockcount by blocksize + + return size, err +} + +// GetBlockdevices gets the available blockdevices for reading from/writing to. +func GetBlockdevices() (devices []Blockdevice, err error) { + block, err := os.ReadDir("/sys/block") + if err != nil { + return + } else if len(block) == 0 { + return devices, fmt.Errorf("no blockdevices found") + } + + for _, device := range block { + var dev Blockdevice + + dev.Path = "/dev/" + device.Name() + + if strings.Contains(device.Name(), "sd") { + dev.Type = "sata" + } else if strings.Contains(device.Name(), "hd") { + dev.Type = "ide" + } else if strings.Contains(device.Name(), "nvme") { + dev.Type = "nvme" + } else if strings.Contains(device.Name(), "dm") { + dev.Type = "dmcrypt" + } else if strings.Contains(device.Name(), "md") { + dev.Type = "raid" + } + + dev.Size, err = getBlockdeviceSize(device.Name()) + if err != nil { + return + } + + devices = append(devices, dev) + } + + return devices, nil +} diff --git a/main.go b/main.go index c946eed..1492471 100644 --- a/main.go +++ b/main.go @@ -1,302 +1,7 @@ package main -import ( - "crypto/sha256" - "fmt" - ac "github.com/JoaoDanielRufino/go-input-autocomplete" - "github.com/briandowns/spinner" - "github.com/fatih/color" - "github.com/schollz/progressbar/v3" - flag "github.com/spf13/pflag" - "io" - "log" - "os" - "runtime" - "strconv" - "strings" - "time" -) - -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() -} - -func GetPath() string { - path, err := ac.Read("[ " + color.YellowString("i") + " ] Please input your image file: ") - if err != nil { - log.Fatal(err) - } - path = strings.TrimSpace(path) - if strings.HasPrefix(path, "~/") { - homedir, err := os.UserHomeDir() - if err != nil { - log.Fatal(err) - } - return homedir + path[1:] - } - if path == "" { - fmt.Println("[", color.RedString("!"), "] No image given, retrying.") - path = GetPath() - } - return path -} - -func GetDest() string { - PrintAvail() - dest, err := ac.Read("[ " + color.YellowString("i") + " ] Please input destination: ") - if err != nil { - log.Fatal(err) - } - dest = strings.TrimSpace(dest) - if strings.HasPrefix(dest, "~/") { - homedir, err := os.UserHomeDir() - if err != nil { - log.Fatal(err) - } - return homedir + dest[1:] - } - if dest == "" { - fmt.Println("[", color.RedString("!"), "] No destination given, retrying.") - dest = GetDest() - } - return dest -} - -func WriteImage(image *os.File, target *os.File, size int64) (int64, error) { - bar := progressbar.NewOptions(int(size), - progressbar.OptionSetWriter(os.Stderr), - progressbar.OptionShowBytes(true), - progressbar.OptionSetWidth(25), - progressbar.OptionSetDescription("Writing image file..."), - progressbar.OptionSetTheme(progressbar.Theme{ - Saucer: "=", - SaucerHead: ">", - SaucerPadding: " ", - BarStart: "[", - BarEnd: "]", - })) - writer := io.MultiWriter(target, bar) - written, err := io.Copy(writer, image) - if err != nil { - return 0, err - } - return written, err -} - -func PrintAvail() { - if runtime.GOOS == "linux" { - block, _ := os.ReadDir("/sys/block") - if len(block) == 0 { - return - } - var targets []string - for _, device := range block { - if strings.HasPrefix(device.Name(), "sd") { - targets = append(targets, device.Name()) - } - if strings.HasPrefix(device.Name(), "nvme") { - targets = append(targets, device.Name()) - } - if strings.HasPrefix(device.Name(), "vd") { - targets = append(targets, device.Name()) - } - } - fmt.Println("Available devices:") - for _, target := range targets { - sizefile, _ := os.Open("/sys/block/" + target + "/size") - sizeread, _ := io.ReadAll(sizefile) - _ = sizefile.Close() - sizestring := strings.ReplaceAll(string(sizeread), "\n", "") - size, _ := strconv.Atoi(sizestring) - size = size * 512 - size = size / 1024 / 1024 / 1024 - - fmt.Print(" * ", "/dev/"+target) - if size > 0 { - fmt.Print(" [", size, "GB]\n") - } else { - fmt.Println("") - } - } - } -} +import "github.com/byReqz/go-etcher/cmd/cetcher" func main() { - s := spinner.New(spinner.CharSets[14], 100*time.Millisecond) - - if input == "" { - if len(flag.Args()) == 0 { - input = GetPath() - } else if len(flag.Args()) > 0 { - input = flag.Args()[0] - } - } - - if device == "" { - if len(flag.Args()) == 0 { - device = GetDest() - } else if len(flag.Args()) > 0 { - if input == flag.Args()[0] && len(flag.Args()) > 1 { - device = flag.Args()[1] - } else if input != flag.Args()[0] && len(flag.Args()) > 0 { - device = flag.Args()[0] - } - } - } - - s.Prefix = "[ " - s.Suffix = " ] Getting file details" - s.Start() - statinput, err := os.Stat(input) - if err != nil { - s.Stop() - fmt.Println("\r[", color.RedString("✘"), "] Getting file details ") - log.Fatal(err) - } - statdevice, err := os.Stat(device) - if err != nil { - s.Stop() - fmt.Println("\r[", color.RedString("✘"), "] Getting file details ") - log.Fatal(err) - } - image, err := os.Open(input) - if err != nil { - s.Stop() - fmt.Println("\r[", color.RedString("✘"), "] Getting file details ") - log.Fatal(err) - } - var inputsize int64 - var inputisblock bool - if statinput.Size() != 0 { - inputsize = statinput.Size() - inputisblock = false - } else { - inputsize, _ = image.Seek(0, io.SeekEnd) - inputisblock = true - _, _ = image.Seek(0, 0) - } - target, err := os.OpenFile(device, os.O_RDWR, 0660) - if err != nil { - s.Stop() - fmt.Println("\r[", color.RedString("✘"), "] Getting file details ") - log.Fatal(err) - } - var targetsize int64 - var targetisblock bool - if statdevice.Size() != 0 { - targetsize = statdevice.Size() - targetisblock = false - } else { - targetsize, err = target.Seek(0, io.SeekEnd) - 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 ") - log.Fatal(err) - } else { - s.Stop() - fmt.Println("\r[", color.GreenString("✓"), "] Getting file details ") - } - inputmb := fmt.Sprint("[", inputsize/1024/1024, "MB]") - devicemb := fmt.Sprint("[", targetsize/1024/1024, "MB]") - var inputblock string - var targetblock string - if inputisblock { - inputblock = "[Blockdevice]" - } else { - inputblock = "[File]" - } - if targetisblock { - targetblock = "[Blockdevice]" - } else { - targetblock = "[File]" - } - fmt.Println("[", color.BlueString("i"), "] Input device/file: "+input, inputmb, inputblock) - fmt.Println("[", color.BlueString("i"), "] Output device/file: "+device, devicemb, targetblock) - if !force { - if inputsize > targetsize { - fmt.Println("[", color.RedString("w"), "]", color.RedString(" Warning:"), "Input file seems to be bigger than the destination!") - } - fmt.Print(color.HiWhiteString("Do you want to continue? [y/N]: ")) - var yesno string - _, _ = fmt.Scanln(&yesno) - yesno = strings.TrimSpace(yesno) - if !(yesno == "y" || yesno == "Y") { - log.Fatal("aborted") - } - } - 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) - } else { - fmt.Println("\r[", color.GreenString("✓"), "] Writing image,", written, "bytes written ") - } - - s.Prefix = "[ " - s.Suffix = " ] Syncing" - s.Start() - 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 ") - log.Fatal(err) - } else { - 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) - } -} + cetcher.Main() +} \ No newline at end of file