From d4bdd9f3853c4dde0dc2ebeb08b1cfea33265933 Mon Sep 17 00:00:00 2001
From: Craig Raw <craigraw@gmail.com>
Date: Tue, 30 Apr 2024 11:54:43 +0200
Subject: [PATCH] support finding signing nodes from provided psbt input
 derivation paths

---
 .../sparrowwallet/drongo/wallet/Keystore.java  | 18 ++++++++++++++++++
 .../sparrowwallet/drongo/wallet/Wallet.java    | 15 +++++++++++++--
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/src/main/java/com/sparrowwallet/drongo/wallet/Keystore.java b/src/main/java/com/sparrowwallet/drongo/wallet/Keystore.java
index 27fb332..19e3fee 100644
--- a/src/main/java/com/sparrowwallet/drongo/wallet/Keystore.java
+++ b/src/main/java/com/sparrowwallet/drongo/wallet/Keystore.java
@@ -242,6 +242,24 @@ public class Keystore extends Persistable {
         return getKeyDerivation().extend(keyPurpose.getPathIndex()).extend(new ChildNumber(keyIndex));
     }
 
+    public ECKey getPubKeyForDerivation(KeyDerivation keyDerivation) {
+        if(keyDerivation != null) {
+            List<ChildNumber> derivation = keyDerivation.getDerivation();
+            if(derivation.size() > this.keyDerivation.getDerivation().size()) {
+                List<ChildNumber> xpubDerivation = derivation.subList(0, this.keyDerivation.getDerivation().size());
+                if(xpubDerivation.equals(this.keyDerivation.getDerivation())) {
+                    derivation = derivation.subList(this.keyDerivation.getDerivation().size(), derivation.size());
+                }
+            }
+
+            if(derivation.size() == 2) {
+                return getPubKey(new WalletNode(KeyDerivation.writePath(derivation)));
+            }
+        }
+
+        return null;
+    }
+
     public boolean isValid() {
         try {
             checkKeystore();
diff --git a/src/main/java/com/sparrowwallet/drongo/wallet/Wallet.java b/src/main/java/com/sparrowwallet/drongo/wallet/Wallet.java
index 8db827a..65f31fe 100644
--- a/src/main/java/com/sparrowwallet/drongo/wallet/Wallet.java
+++ b/src/main/java/com/sparrowwallet/drongo/wallet/Wallet.java
@@ -1490,8 +1490,8 @@ public class Wallet extends Persistable implements Comparable<Wallet> {
         return signingNodes;
     }
 
-    public List<Keystore> getSigningKeystores(PSBT psbt) {
-        List<Keystore> signingKeystores = new ArrayList<>();
+    public Collection<Keystore> getSigningKeystores(PSBT psbt) {
+        Set<Keystore> signingKeystores = new LinkedHashSet<>();
 
         for(Map.Entry<ExtendedKey, KeyDerivation> entry : psbt.getExtendedPublicKeys().entrySet()) {
             for(Keystore keystore : getKeystores()) {
@@ -1501,6 +1501,17 @@ public class Wallet extends Persistable implements Comparable<Wallet> {
             }
         }
 
+        for(PSBTInput psbtInput : psbt.getPsbtInputs()) {
+            for(Map.Entry<ECKey, KeyDerivation> entry : psbtInput.getDerivedPublicKeys().entrySet()) {
+                for(Keystore keystore : getKeystores().stream().filter(k -> !signingKeystores.contains(k)).toList()) {
+                    ECKey derivedKey = keystore.getPubKeyForDerivation(entry.getValue());
+                    if(derivedKey != null && Arrays.equals(entry.getKey().getPubKey(), derivedKey.getPubKey())) {
+                        signingKeystores.add(keystore);
+                    }
+                }
+            }
+        }
+
         return signingKeystores;
     }