restructure packages and start picking apart functions

This commit is contained in:
Nils 2023-10-08 11:57:13 +02:00
parent d0f0957932
commit 0625418d9d
Signed by: byreqz
GPG Key ID: 396A62D7D436749E
4 changed files with 365 additions and 298 deletions

283
cmd/cetcher/main.go Normal file
View File

@ -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)
}
}

11
etcher/etcher.go Normal file
View File

@ -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 {}

68
etcher/etcher_linux.go Normal file
View File

@ -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
}

301
main.go
View File

@ -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()
}