From c87945d752ddf9ed5ac218db929df8d881b998bb Mon Sep 17 00:00:00 2001
From: Michele Balistreri <michele@bitgamma.com>
Date: Fri, 29 Nov 2024 11:22:11 +0900
Subject: [PATCH] support exporting xpub keys

---
 command_set.go        | 27 +++++++++++++++++++--------
 types/exported_key.go | 39 ++++++++++++++++++++++++++++++---------
 2 files changed, 49 insertions(+), 17 deletions(-)

diff --git a/command_set.go b/command_set.go
index 095d9c6..9aee963 100644
--- a/command_set.go
+++ b/command_set.go
@@ -316,6 +316,23 @@ func (cs *CommandSet) DeriveKey(path string) error {
 }
 
 func (cs *CommandSet) ExportKey(derive bool, makeCurrent bool, onlyPublic bool, path string) ([]byte, []byte, error) {
+	var p2 uint8
+	if onlyPublic {
+		p2 = P2ExportKeyPublicOnly
+	} else {
+		p2 = P2ExportKeyPrivateAndPublic
+	}
+
+	key, err := cs.ExportKeyExtended(derive, makeCurrent, p2, path)
+
+	if err != nil {
+		return nil, nil, err
+	}
+
+	return key.PrivKey(), key.PubKey(), err
+}
+
+func (cs *CommandSet) ExportKeyExtended(derive bool, makeCurrent bool, p2 uint8, path string) (*types.ExportedKey, error) {
 	var p1 uint8
 	if !derive {
 		p1 = P1ExportKeyCurrent
@@ -324,22 +341,16 @@ func (cs *CommandSet) ExportKey(derive bool, makeCurrent bool, onlyPublic bool,
 	} else {
 		p1 = P1ExportKeyDeriveAndMakeCurrent
 	}
-	var p2 uint8
-	if onlyPublic {
-		p2 = P2ExportKeyPublicOnly
-	} else {
-		p2 = P2ExportKeyPrivateAndPublic
-	}
 
 	cmd, err := NewCommandExportKey(p1, p2, path)
 	if err != nil {
-		return nil, nil, err
+		return nil, err
 	}
 
 	resp, err := cs.sc.Send(cmd)
 	err = cs.checkOK(resp, err)
 	if err != nil {
-		return nil, nil, err
+		return nil, err
 	}
 
 	return types.ParseExportKeyResponse(resp.Data)
diff --git a/types/exported_key.go b/types/exported_key.go
index cd8a472..ae856ea 100644
--- a/types/exported_key.go
+++ b/types/exported_key.go
@@ -8,29 +8,50 @@ import (
 )
 
 var (
-	TagExportKeyTemplate = uint8(0xA1)
-	TagExportKeyPublic   = uint8(0x81)
+	TagExportKeyTemplate    = apdu.Tag{0xA1}
+	TagExportKeyPublic      = apdu.Tag{0x80}
+	TagExportKeyPrivate     = apdu.Tag{0x81}
+	TagExportKeyPublicChain = apdu.Tag{0x82}
 )
 
-func ParseExportKeyResponse(data []byte) ([]byte, []byte, error) {
-	tpl, err := apdu.FindTag(data, apdu.Tag{0xA1})
+type ExportedKey struct {
+	pubKey    []byte
+	privKey   []byte
+	chainCode []byte
+}
+
+func (k *ExportedKey) PubKey() []byte {
+	return k.pubKey
+}
+
+func (k *ExportedKey) PrivKey() []byte {
+	return k.privKey
+}
+
+func (k *ExportedKey) ChainCode() []byte {
+	return k.chainCode
+}
+
+func ParseExportKeyResponse(data []byte) (*ExportedKey, error) {
+	tpl, err := apdu.FindTag(data, TagExportKeyTemplate)
 	if err != nil {
-		return nil, nil, err
+		return nil, err
 	}
 
-	pubKey := tryFindTag(tpl, apdu.Tag{0x80})
-	privKey := tryFindTag(tpl, apdu.Tag{0x81})
+	pubKey := tryFindTag(tpl, TagExportKeyPublic)
+	privKey := tryFindTag(tpl, TagExportKeyPrivate)
+	chainCode := tryFindTag(tpl, TagExportKeyPublicChain)
 
 	if len(pubKey) == 0 && len(privKey) > 0 {
 		ecdsaKey, err := ethcrypto.HexToECDSA(fmt.Sprintf("%x", privKey))
 		if err != nil {
-			return nil, nil, err
+			return nil, err
 		}
 
 		pubKey = ethcrypto.FromECDSAPub(&ecdsaKey.PublicKey)
 	}
 
-	return privKey, pubKey, nil
+	return &ExportedKey{pubKey, privKey, chainCode}, nil
 }
 
 func tryFindTag(tpl []byte, tags ...apdu.Tag) []byte {