first commit

8 files changed, 478 insertions, 0 deletions

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..b5ed5dc
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1 @@
+wally-cli
diff --git a/README.md b/README.md
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/README.md
diff --git a/dfu.go b/dfu.go
new file mode 100644
index 0000000..52b3283
--- /dev/null
+++ b/dfu.go
@@ -0,0 +1,223 @@
+package main
+
+import (
+	"errors"
+	"fmt"
+	"github.com/google/gousb"
+	"io/ioutil"
+	"log"
+	"time"
+)
+
+type status struct {
+	bStatus       string
+	bwPollTimeout int
+	bState        string
+	iString       string
+}
+
+func dfuCommand(dev *gousb.Device, addr int, command int, status *status) (err error) {
+	var buf []byte
+	if command == setAddress {
+		buf = make([]byte, 5)
+		buf[0] = 0x21
+		buf[1] = byte(addr & 0xff)
+		buf[2] = byte((addr >> 8) & 0xff)
+		buf[3] = byte((addr >> 16) & 0xff)
+		buf[4] = byte((addr >> 24) & 0xff)
+	}
+	if command == eraseAddress {
+		buf = make([]byte, 5)
+		buf[0] = 0x41
+		buf[1] = byte(addr & 0xff)
+		buf[2] = byte((addr >> 8) & 0xff)
+		buf[3] = byte((addr >> 16) & 0xff)
+		buf[4] = byte((addr >> 24) & 0xff)
+	}
+	if command == eraseFlash {
+		buf = make([]byte, 1)
+		buf[0] = 0x41
+	}
+
+	_, err = dev.Control(33, 1, 0, 0, buf)
+
+	err = dfuPollTimeout(dev, status)
+
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func dfuPollTimeout(dev *gousb.Device, status *status) (err error) {
+	for i := 0; i < 3; i++ {
+		err = dfuGetStatus(dev, status)
+		time.Sleep(time.Duration(status.bwPollTimeout) * time.Millisecond)
+	}
+	return err
+}
+
+func dfuGetStatus(dev *gousb.Device, status *status) (err error) {
+	buf := make([]byte, 6)
+	stat, err := dev.Control(161, 3, 0, 0, buf)
+	if err != nil {
+		return err
+	}
+	if stat == 6 {
+		status.bStatus = string(buf[0])
+		status.bwPollTimeout = int((0xff & buf[3] << 16) | (0xff & buf[2]) | 0xff&buf[1])
+		status.bState = string(buf[4])
+		status.iString = string(buf[5])
+	}
+	return err
+}
+
+func dfuClearStatus(dev *gousb.Device) (err error) {
+	_, err = dev.Control(33, 4, 2, 0, nil)
+	return err
+}
+
+func dfuReboot(dev *gousb.Device, status *status) (err error) {
+	err = dfuPollTimeout(dev, status)
+	_, err = dev.Control(33, 1, 2, 0, nil)
+	time.Sleep(1000 * time.Millisecond)
+	err = dfuGetStatus(dev, status)
+	return err
+}
+
+func extractSuffix(fileData []byte) (hasSuffix bool, data []byte, err error) {
+
+	fileSize := len(fileData)
+
+	suffix := fileData[fileSize-dfuSuffixLength : fileSize]
+	d := string(suffix[10])
+	f := string(suffix[9])
+	u := string(suffix[8])
+
+	if d == "D" && f == "F" && u == "U" {
+		vid := int((suffix[5] << 8) + suffix[4])
+		pid := int((suffix[3] << 8) + suffix[2])
+		if vid != dfuSuffixVendorID || pid != dfuSuffixProductID {
+			message := fmt.Sprintf("Invalid vendor or product id, expected %#x:%#x got %#x:%#x", dfuSuffixVendorID, dfuSuffixProductID, vid, pid)
+			err = errors.New(message)
+			return true, fileData, err
+
+		}
+
+		return true, fileData[0 : fileSize-dfuSuffixLength], nil
+	}
+
+	return false, fileData, nil
+}
+
+func dfuFlash(firmwarePath string, s *state) {
+	dfuStatus := status{}
+	fileData, err := ioutil.ReadFile(firmwarePath)
+	if err != nil {
+		message := fmt.Sprintf("Error while opening firmware: %s", err)
+		log.Fatal(message)
+		return
+	}
+
+	_, firmwareData, err := extractSuffix(fileData)
+	if err != nil {
+		message := fmt.Sprintf("Error while extracting DFU Suffix: %s", err)
+		log.Fatal(message)
+		return
+	}
+
+	ctx := gousb.NewContext()
+	defer ctx.Close()
+	var dev *gousb.Device
+
+	// Get the list of device that match TMK's vendor id
+	for {
+		devs, _ := ctx.OpenDevices(func(desc *gousb.DeviceDesc) bool {
+			if desc.Vendor == gousb.ID(dfuVendorID) && desc.Product == gousb.ID(dfuProductID) {
+				return true
+			}
+			return false
+		})
+
+		defer func() {
+			for _, d := range devs {
+				d.Close()
+			}
+		}()
+
+		if len(devs) > 0 {
+			dev = devs[0]
+			break
+		}
+		time.Sleep(1 * time.Second)
+	}
+
+	dev.SetAutoDetach(true)
+
+	dev.ControlTimeout = 5 * time.Second
+
+	cfg, err := dev.Config(1)
+	if err != nil {
+		message := fmt.Sprintf("Error while claiming the usb interface: %s", err)
+		log.Fatal(message)
+	}
+	defer cfg.Close()
+
+	fileSize := len(firmwareData)
+	s.total = fileSize
+
+	err = dfuClearStatus(dev)
+	if err != nil {
+		message := fmt.Sprintf("Error while clearing the device status: %s", err)
+		log.Fatal(message)
+	}
+
+	s.step = 1
+
+	err = dfuCommand(dev, 0, eraseFlash, &dfuStatus)
+	if err != nil {
+		message := fmt.Sprintf("Error while erasing flash:", err)
+		log.Fatal(message)
+		return
+	}
+
+	for page := 0; page < fileSize; page += planckBlockSize {
+		addr := planckStartAddress + page
+		chunckSize := planckBlockSize
+
+		if page+chunckSize > fileSize {
+			chunckSize = fileSize - page
+		}
+
+		err = dfuCommand(dev, addr, eraseAddress, &dfuStatus)
+		if err != nil {
+			message := fmt.Sprintf("Error while sending the erase address command: %s", err)
+			log.Fatal(message)
+		}
+		err = dfuCommand(dev, addr, setAddress, &dfuStatus)
+		if err != nil {
+			message := fmt.Sprintf("Error while sending the set address command: %s", err)
+			log.Fatal(message)
+		}
+
+		buf := firmwareData[page : page+chunckSize]
+		bytes, err := dev.Control(33, 1, 2, 0, buf)
+
+		if err != nil {
+			message := fmt.Sprintf("Error while sending firmware bytes: %s", err)
+			log.Fatal(message)
+		}
+
+		s.sent += bytes
+	}
+
+	err = dfuReboot(dev, &dfuStatus)
+	if err != nil {
+		message := fmt.Sprintf("Error while rebooting device: %s", err)
+		log.Fatal(message)
+		return
+	}
+
+	s.step = 2
+}
diff --git a/go.mod b/go.mod
new file mode 100644
index 0000000..702d099
--- /dev/null
+++ b/go.mod
@@ -0,0 +1,13 @@
+module wally-cli
+
+go 1.14
+
+require (
+	github.com/caarlos0/spin v1.1.0
+	github.com/google/gousb v2.1.0+incompatible
+	github.com/logrusorgru/aurora v2.0.3+incompatible
+	github.com/marcinbor85/gohex v0.0.0-20200531163658-baab2527a9a2
+	github.com/mattn/go-runewidth v0.0.9 // indirect
+	golang.org/x/sys v0.0.0-20200923182605-d9f96fdee20d // indirect
+	gopkg.in/cheggaaa/pb.v1 v1.0.28
+)
diff --git a/go.sum b/go.sum
new file mode 100644
index 0000000..3ed918e
--- /dev/null
+++ b/go.sum
@@ -0,0 +1,15 @@
+github.com/caarlos0/spin v1.1.0 h1:EjsfGbZJejib25BPnDqf7iL2z9RUna7refvUf+AN9UE=
+github.com/caarlos0/spin v1.1.0/go.mod h1:HOC4pUvfhjXR2yDt+sEY9dRc2m4CCaK5z5oQYAbzXSA=
+github.com/google/gousb v1.1.0 h1:s/970WE1z968MC+dtWbuxDHCcx9kwANQo6UcZtfTfx0=
+github.com/google/gousb v2.1.0+incompatible h1:ApzMDjF3FeO219QwWybJxYfFhXQzPLOEy0o+w9k5DNI=
+github.com/google/gousb v2.1.0+incompatible/go.mod h1:Tl4HdAs1ThE3gECkNwz+1MWicX6FXddhJEw7L8jRDiI=
+github.com/logrusorgru/aurora v2.0.3+incompatible h1:tOpm7WcpBTn4fjmVfgpQq0EfczGlG91VSDkswnjF5A8=
+github.com/logrusorgru/aurora v2.0.3+incompatible/go.mod h1:7rIyQOR62GCctdiQpZ/zOJlFyk6y+94wXzv6RNZgaR4=
+github.com/marcinbor85/gohex v0.0.0-20200531163658-baab2527a9a2 h1:n7R8fUwWZUB2XtyzBNsYNNm9/XgOBj6pvLi7GLMCHtM=
+github.com/marcinbor85/gohex v0.0.0-20200531163658-baab2527a9a2/go.mod h1:Pb6XcsXyropB9LNHhnqaknG/vEwYztLkQzVCHv8sQ3M=
+github.com/mattn/go-runewidth v0.0.9 h1:Lm995f3rfxdpd6TSmuVCHVb/QhupuXlYr8sCI/QdE+0=
+github.com/mattn/go-runewidth v0.0.9/go.mod h1:H031xJmbD/WCDINGzjvQ9THkh0rPKHF+m2gUSrubnMI=
+golang.org/x/sys v0.0.0-20200923182605-d9f96fdee20d h1:L/IKR6COd7ubZrs2oTnTi73IhgqJ71c9s80WsQnh0Es=
+golang.org/x/sys v0.0.0-20200923182605-d9f96fdee20d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+gopkg.in/cheggaaa/pb.v1 v1.0.28 h1:n1tBJnnK2r7g9OW2btFH91V92STTUevLXYFb8gy9EMk=
+gopkg.in/cheggaaa/pb.v1 v1.0.28/go.mod h1:V/YB90LKu/1FcN3WVnfiiE5oMCibMjukxqG/qStrOgw=
diff --git a/main.go b/main.go
new file mode 100644
index 0000000..a2df1fe
--- /dev/null
+++ b/main.go
@@ -0,0 +1,78 @@
+package main
+
+import (
+	"fmt"
+	"github.com/caarlos0/spin"
+	"github.com/logrusorgru/aurora"
+	"gopkg.in/cheggaaa/pb.v1"
+	"os"
+	"path/filepath"
+	"time"
+)
+
+var appVersion = "2.0.0"
+
+type state struct {
+	step  int
+	total int
+	sent  int
+}
+
+func main() {
+	var args = os.Args[1:]
+	s := state{step: 0, total: 0, sent: 0}
+
+	if len(args) != 1 {
+		fmt.Println(aurora.Blue("Usage: wally-cli <firmware file>"))
+		os.Exit(0)
+	}
+
+	if args[0] == "--version" {
+		appVersion := fmt.Sprintf("wally-cli v%s", appVersion)
+		fmt.Println(aurora.Blue(appVersion))
+		os.Exit(0)
+	}
+
+	path := args[0]
+	if _, err := os.Stat(path); os.IsNotExist(err) {
+		fmt.Println(aurora.Red("The file path you specified does not exist"))
+		os.Exit(1)
+	}
+
+	extension := filepath.Ext(path)
+	if extension != ".bin" && extension != ".hex" {
+		fmt.Println(aurora.Red("Please provide a valid firmware file: a"), aurora.Red(aurora.Underline(".hex")), aurora.Red("file (ErgoDox EZ) or a"), aurora.Red(aurora.Underline(".bin")), aurora.Red("file (Moonlander / Planck EZ)"))
+		os.Exit(1)
+	}
+
+	spinner := spin.New("%s Press the reset button of your keyboard.")
+	spinner.Start()
+	spinnerStopped := false
+
+	var progress *pb.ProgressBar
+	progressStarted := false
+
+	if extension == ".bin" {
+		go dfuFlash(path, &s)
+	}
+	if extension == ".hex" {
+		go teensyFlash(path, &s)
+	}
+
+	for s.step != 2 {
+		time.Sleep(500 * time.Millisecond)
+		if s.step > 0 {
+			if spinnerStopped == false {
+				spinner.Stop()
+				spinnerStopped = true
+			}
+			if progressStarted == false {
+				progressStarted = true
+				progress = pb.StartNew(s.total)
+			}
+			progress.Set(s.sent)
+		}
+	}
+	progress.Finish()
+	fmt.Println(aurora.Green("Your keyboard was successfully flashed and rebooted. Enjoy the new firmware!"))
+}
diff --git a/teensy.go b/teensy.go
new file mode 100644
index 0000000..02c9908
--- /dev/null
+++ b/teensy.go
@@ -0,0 +1,114 @@
+package main
+
+import (
+	"fmt"
+	"github.com/google/gousb"
+	"github.com/marcinbor85/gohex"
+	"log"
+	"os"
+	"time"
+)
+
+// TeensyFlash: Flashes Teensy boards.
+// It opens the firmware file at the provided path, checks it's integrity, wait for the keyboard to be in Flash mode, flashes it and reboots the board.
+func teensyFlash(firmwarePath string, s *state) {
+	file, err := os.Open(firmwarePath)
+	if err != nil {
+		message := fmt.Sprintf("Error while opening firmware: %s", err)
+		log.Fatal(message)
+		return
+	}
+	defer file.Close()
+
+	s.total = ergodoxCodeSize
+
+	firmware := gohex.NewMemory()
+	err = firmware.ParseIntelHex(file)
+	if err != nil {
+		message := fmt.Sprintf("Error while parsing firmware: %s", err)
+		log.Fatal(message)
+		return
+	}
+
+	ctx := gousb.NewContext()
+	defer ctx.Close()
+	var dev *gousb.Device
+
+	// Loop until a keyboard is ready to flash
+	for {
+		devs, _ := ctx.OpenDevices(func(desc *gousb.DeviceDesc) bool {
+			if desc.Vendor == gousb.ID(halfKayVendorID) && desc.Product == gousb.ID(halfKayProductID) {
+				return true
+			}
+			return false
+		})
+
+		defer func() {
+			for _, d := range devs {
+				d.Close()
+			}
+		}()
+
+		if len(devs) > 0 {
+			dev = devs[0]
+			break
+		}
+		time.Sleep(1 * time.Second)
+	}
+
+	// Detach keyboard from the kernel
+	dev.SetAutoDetach(true)
+
+	// Claim usb device
+	cfg, err := dev.Config(1)
+	defer cfg.Close()
+	if err != nil {
+		message := fmt.Sprintf("Error while claiming the usb interface: %s", err)
+		log.Fatal(message)
+		return
+	}
+
+	s.step = 1
+
+	// Loop on the firmware data and program
+	var addr uint32
+	for addr = 0; addr < ergodoxCodeSize; addr += ergodoxBlockSize {
+		// set a longer timeout when writing the first block
+		if addr == 0 {
+			dev.ControlTimeout = 5 * time.Second
+		} else {
+			dev.ControlTimeout = 500 * time.Millisecond
+		}
+		// Prepare and write a firmware block
+		// https://www.pjrc.com/teensy/halfkay_protocol.html
+		buf := make([]byte, ergodoxBlockSize+2)
+		buf[0] = byte(addr & 255)
+		buf[1] = byte((addr >> 8) & 255)
+		block := firmware.ToBinary(addr, ergodoxBlockSize, 255)
+		for index := range block {
+			buf[index+2] = block[index]
+		}
+
+		bytes, err := dev.Control(0x21, 9, 0x0200, 0, buf)
+		if err != nil {
+			message := fmt.Sprintf("Error while sending firmware bytes: %s", err)
+			log.Fatal(message)
+			return
+		}
+
+		s.sent += bytes
+	}
+
+	buf := make([]byte, ergodoxBlockSize+2)
+	buf[0] = byte(0xFF)
+	buf[1] = byte(0xFF)
+	buf[2] = byte(0xFF)
+	_, err = dev.Control(0x21, 9, 0x0200, 0, buf)
+
+	if err != nil {
+		message := fmt.Sprintf("Error while rebooting device: %s", err)
+		log.Fatal(message)
+		return
+	}
+	s.step = 2
+}
diff --git a/usb.go b/usb.go
new file mode 100644
index 0000000..e1cbe0f
--- /dev/null
+++ b/usb.go
@@ -0,0 +1,34 @@
+package main
+
+type device struct {
+	model int // 0 - planck // 1 - ergodox // 2 - moonlander
+	bus   int
+	port  int
+}
+
+const (
+	vendorID1    int = 0xFEED
+	vendorID2    int = 0x3297
+	planckID     int = 0x6060
+	ergodoxID    int = 0x1307
+	moonlanderID int = 0x1969
+
+	dfuSuffixVendorID  int = 0x83
+	dfuSuffixProductID int = 0x11
+	dfuVendorID        int = 0x0483
+	dfuProductID       int = 0xdf11
+
+	halfKayVendorID  int = 0x16C0
+	halfKayProductID int = 0x0478
+
+	ergodoxMaxMemorySize = 0x100000
+	ergodoxCodeSize      = 32256
+	ergodoxBlockSize     = 128
+
+	dfuSuffixLength    = 16
+	planckBlockSize    = 2048
+	planckStartAddress = 0x08000000
+	setAddress         = 0
+	eraseAddress       = 1
+	eraseFlash         = 2
+)