feature: add flash subcommand and move existing functionality to keys subcommand

Signed-off-by: deadprogram <[email protected]>

Author: deadprogram

README.md

@@ -72,18 +72,18 @@ go install github.com/hybridgroup/go-haystack
 1. Generate keys for a device
 
 ```shell
-go-haystack -name=DEVICENAME
+go-haystack keys DEVICENAME
 ```
 
 The keys will be saved in a file named `DEVICENAME.keys` and the configuration file for Haystack will be saved in `DEVICENAME.json`. Replace "DEVICENAME" with whatever you want to name the actual device.
 
 
-2. Flash the hardware with the target name of your device and those keys
+2. Flash the hardware with the TinyGo target and the name of your device.
 
 For example:
 
 ```shell
-./flash.sh nano-rp2040 DEVICENAME.keys
+go-haystack flash nano-rp2040 DEVICENAME
 ```
 
 This will use TinyGo to compile the firmware using your keys, and then flash it to the device. See [https://tinygo.org/getting-started/overview/](https://tinygo.org/getting-started/overview/) for more information about TinyGo.

keys.go

@@ -6,22 +6,25 @@ import (
 	"crypto/rand"
 	"crypto/sha256"
 	"encoding/base64"
+	"errors"
 	"strings"
 )
 
-func generateKey() (string, string, string) {
+var errInvalidHash = errors.New("Hash contains '/'")
+
+func generateKey() (string, string, string, error) {
 	// Generate ECDSA private key using P-224 curve
 	pk, err := ecdsa.GenerateKey(elliptic.P224(), rand.Reader)
 	if err != nil {
-		panic(err)
+		return "", "", "", err
 	}
 
 	// Extract the raw private key bytes
 	privateKeyBytes := pk.D.Bytes()
 
 	// Ensure the private key is 28 bytes long (P-224 curve)
 	if len(privateKeyBytes) != 28 {
-		panic("Private key is not 28 bytes long")
+		return "", "", "", errors.New("Private key is not 28 bytes long")
 	}
 
 	// Encode the raw private key to Base64
@@ -39,8 +42,8 @@ func generateKey() (string, string, string) {
 
 	// make sure not '/' in the base64 string
 	if strings.Contains(hashBase64, "/") {
-		panic("key generation failed because there was a / in the b64 of the hashed pubkey. try again.")
+		return "", "", "", errInvalidHash
 	}
 
-	return privateKeyBase64, publicKeyBase64, hashBase64
+	return privateKeyBase64, publicKeyBase64, hashBase64, nil
 }

main.go

@@ -1,21 +1,55 @@
 package main
 
 import (
+	"errors"
 	"flag"
 	"fmt"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strings"
 )
 
 func main() {
-	nameFlag := flag.String("name", "", "name of the device to be added")
-	verboseFlag := flag.Bool("verbose", false, "enable verbose mode")
-
+	verboseFlag := flag.Bool("v", false, "enable verbose mode")
 	flag.Parse()
-	if *nameFlag == "" {
-		fmt.Println("Please provide a device name using the -name flag")
+
+	args := flag.Args()
+	if len(args) < 1 {
+		fmt.Println("either 'keys' or 'flash' subcommand is required.", len(args))
 		return
 	}
 
-	priv, pub, hash := generateKey()
+	switch args[0] {
+	case "keys":
+		if len(args) < 2 {
+			fmt.Println("Please provide a device name")
+			return
+		}
+		if err := generateKeys(args[1], verboseFlag); err != nil {
+			fmt.Println("failed to generate keys:", err)
+		}
+	case "flash":
+		if len(args) < 3 {
+			fmt.Println("Please provide a device target and name")
+			return
+		}
+		if err := flashDevice(args[1], args[2], verboseFlag); err != nil {
+			fmt.Println("failed to flash device:", err)
+		}
+	default:
+		fmt.Println("either 'keys' or 'flash' subcommand is required.")
+		return
+	}
+}
+
+func generateKeys(name string, verboseFlag *bool) error {
+	// TODO: check if overwriting keys
+
+	priv, pub, hash, err := generateKey()
+	if err != nil {
+		return err
+	}
 
 	// Print the keys and hash
 	if *verboseFlag {
@@ -25,8 +59,58 @@ func main() {
 	}
 
 	// save keys file
-	saveKeys(*nameFlag, priv, pub, hash)
+	if err := saveKeys(name, priv, pub, hash); err != nil {
+		return err
+	}
 
 	// save device file
-	saveDevice(*nameFlag, priv)
+	return saveDevice(name, priv)
+}
+
+func flashDevice(name string, target string, verboseFlag *bool) error {
+	key, err := readKey(name)
+	if err != nil {
+		return err
+	}
+
+	pwd := os.Getenv("PWD")
+	pth := filepath.Join(pwd, "firmware")
+	if err := os.Chdir(pth); err != nil {
+		panic(err)
+	}
+	defer os.Chdir(pwd)
+
+	keyVal := fmt.Sprintf("-X main.AdvertisingKey='%s'", key)
+	if *verboseFlag {
+		fmt.Println("tinygo", "flash", "-target", target, "-ldflags", keyVal, ".")
+	}
+
+	cmd := exec.Command("tinygo", "flash", "-target", target, "-ldflags", keyVal, ".")
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+	return cmd.Run()
+}
+
+func readKey(name string) (string, error) {
+	f, err := os.Open(name + ".keys")
+	if err != nil {
+		return "", err
+	}
+	defer f.Close()
+
+	b := make([]byte, 1024)
+	n, err := f.Read(b)
+	if err != nil {
+		return "", err
+	}
+
+	lines := strings.Split(string(b[:n]), "\n")
+	for _, line := range lines {
+		if strings.Contains(line, "Advertisement key") {
+			s := strings.Split(line, ":")
+			return strings.TrimLeft(s[1], " "), nil
+		}
+	}
+
+	return "", errors.New("key not found")
 }

save.go

@@ -8,17 +8,19 @@ import (
 	"text/template"
 )
 
-func saveKeys(name string, priv string, pub string, hash string) {
+func saveKeys(name string, priv string, pub string, hash string) error {
 	f, err := os.Create(name + ".keys")
 	if err != nil {
-		panic(err)
+		return err
 	}
 
 	defer f.Close()
 
 	f.Write([]byte(fmt.Sprintf("Private key: %s\n", priv)))
 	f.Write([]byte(fmt.Sprintf("Advertisement key: %s\n", pub)))
 	f.Write([]byte(fmt.Sprintf("Hashed adv key: %s\n", hash)))
+
+	return nil
 }
 
 const deviceTemplate = `[
@@ -42,15 +44,15 @@ const deviceTemplate = `[
 ]
 `
 
-func saveDevice(name string, priv string) {
+func saveDevice(name string, priv string) error {
 	t, err := template.New("device").Parse(deviceTemplate)
 	if err != nil {
-		panic(err)
+		return err
 	}
 
 	f, err := os.Create(name + ".json")
 	if err != nil {
-		panic(err)
+		return err
 	}
 
 	defer f.Close()
@@ -61,8 +63,10 @@ func saveDevice(name string, priv string) {
 		"PrivateKey": priv,
 	})
 	if err != nil {
-		panic(err)
+		return err
 	}
+
+	return nil
 }
 
 // Returns an int >= min, < max