Merge pull request #13 from ZondaX/david-yan/list-ledger-devices

David yan/list ledger devices

Makefile

@@ -0,0 +1,2 @@
+test:
+	go test ledger_test.go ledger.go apduWrapper.go

README.md

@@ -8,7 +8,7 @@ It handles USB (HID) communication and APDU encapsulation.
 
 Linux, OSX and Windows are supported.
 
-# Building
-```
+## Building
+```bash
 go build
 ```

ledger.go

@@ -39,7 +39,7 @@ type Ledger struct {
 	Logging     bool
 }
 
-func NewLedger(dev *hid.Device) *Ledger {
+func newLedger(dev *hid.Device) *Ledger {
 	return &Ledger{
 		device:  *dev,
 		Logging: false,
@@ -67,25 +67,52 @@ func ListDevices() {
 	}
 }
 
+func isLedgerDevice(d hid.DeviceInfo) bool {
+	deviceFound := d.UsagePage == UsagePageLedgerNanoS
+	// Workarounds for possible empty usage pages
+	return deviceFound ||
+		(d.Product == "Nano S" && d.Interface == 0) ||
+		(d.Product == "Nano X" && d.Interface == 0)
+}
+
+func CountLedgerDevices() uint {
+	devices := hid.Enumerate(0, 0)
+
+	count := uint(0)
+	for _, d := range devices {
+		if isLedgerDevice(d) {
+			count++
+		}
+	}
+
+	return count
+}
+
 func FindLedger() (*Ledger, error) {
+	return GetLedger(0)
+}
+
+func GetLedger(requiredIndex uint) (*Ledger, error) {
 	devices := hid.Enumerate(VendorLedger, 0)
 
+	currentIndex := uint(0)
 	for _, d := range devices {
-		deviceFound := d.UsagePage == UsagePageLedgerNanoS
-		// Workarounds for possible empty usage pages
-		deviceFound = deviceFound ||
-			(d.Product == "Nano S" && d.Interface == 0) ||
-			(d.Product == "Nano X" && d.Interface == 0)
-
-		if deviceFound {
-			device, err := d.Open()
-			if err == nil {
-				return NewLedger(device), nil
+		if isLedgerDevice(d) {
+			if currentIndex == requiredIndex {
+				device, err := d.Open()
+				if err != nil {
+					return nil, err
+				}
+				return newLedger(device), nil
+			}
+			currentIndex++
+			if currentIndex > requiredIndex {
+				break
 			}
 		}
 	}
 
-	return nil, errors.New("no ledger connected")
+	return nil, fmt.Errorf("Ledger device (idx %d) not found", requiredIndex)
 }
 
 func ErrorMessage(errorCode uint16) string {

ledger_test.go

@@ -19,24 +19,46 @@
 package ledger_go
 
 import (
-	"encoding/hex"
 	"fmt"
 	"github.com/stretchr/testify/assert"
 	"github.com/zondax/hid"
+	"sync"
 	"testing"
 )
 
+var mux sync.Mutex
+
 func Test_ThereAreDevices(t *testing.T) {
+	mux.Lock()
+	defer mux.Unlock()
+
 	devices := hid.Enumerate(0, 0)
 	assert.NotEqual(t, 0, len(devices))
 }
 
 func Test_ListDevices(t *testing.T) {
+	mux.Lock()
+	defer mux.Unlock()
+
 	ListDevices()
 }
 
+func Test_CountLedgerDevices(t *testing.T) {
+	mux.Lock()
+	defer mux.Unlock()
+
+	count := CountLedgerDevices()
+	println(count)
+	assert.True(t, count > 0)
+}
+
 func Test_FindLedger(t *testing.T) {
+	mux.Lock()
+	defer mux.Unlock()
+
 	ledger, err := FindLedger()
+	defer ledger.Close()
+
 	if err != nil {
 		fmt.Println("\n*********************************")
 		fmt.Println("Did you enter the password??")
@@ -46,32 +68,29 @@ func Test_FindLedger(t *testing.T) {
 	assert.NotNil(t, ledger)
 }
 
-func Test_BasicExchange(t *testing.T) {
-	ledger, err := FindLedger()
+func Test_GetLedger(t *testing.T) {
+	mux.Lock()
+	defer mux.Unlock()
+
+	ledger, err := GetLedger(1)
+	defer ledger.Close()
+
 	if err != nil {
 		fmt.Println("\n*********************************")
 		fmt.Println("Did you enter the password??")
 		fmt.Println("*********************************")
 		t.Fatalf("Error: %s", err.Error())
 	}
 	assert.NotNil(t, ledger)
-
-	message := []byte{0x55, 0, 0, 0, 0}
-
-	for i := 0; i < 10; i++ {
-		response, err := ledger.Exchange(message)
-
-		if err != nil {
-			fmt.Printf("iteration %d\n", i)
-			t.Fatalf("Error: %s", err.Error())
-		}
-
-		assert.Equal(t, 5, len(response))
-	}
 }
 
-func Test_LongExchange(t *testing.T) {
+func Test_BasicExchange(t *testing.T) {
+	mux.Lock()
+	defer mux.Unlock()
+
 	ledger, err := FindLedger()
+	defer ledger.Close()
+
 	if err != nil {
 		fmt.Println("\n*********************************")
 		fmt.Println("Did you enter the password??")
@@ -80,20 +99,17 @@ func Test_LongExchange(t *testing.T) {
 	}
 	assert.NotNil(t, ledger)
 
-	path := "052c000080760000800000008000000000000000000000000000000000000000000000000000000000"
-	pathBytes, err := hex.DecodeString(path)
-	if err != nil {
-		t.Fatalf("invalid path in test")
-	}
+	// Call app info (this should work in main menu and many apps)
+	message := []byte{0xB0, 0x01, 0, 0, 0}
 
-	header := []byte{0x55, 1, 0, 0, byte(len(pathBytes))}
-	message := append(header, pathBytes...)
+	for i := 0; i < 10; i++ {
+		response, err := ledger.Exchange(message)
 
-	response, err := ledger.Exchange(message)
+		if err != nil {
+			fmt.Printf("iteration %d\n", i)
+			t.Fatalf("Error: %s", err.Error())
+		}
 
-	if err != nil {
-		t.Fatalf("Error: %s", err.Error())
+		assert.Equal(t, 15, len(response))
 	}
-
-	assert.Equal(t, 65, len(response))
 }