From f6dcdb6d26c3b40ae1c0f7502b3e526aa8959564 Mon Sep 17 00:00:00 2001
From: Craig Raw <craigraw@gmail.com>
Date: Mon, 20 Jul 2020 17:04:39 +0200
Subject: [PATCH] determine signing status from psbt

---
 .../sparrowwallet/drongo/wallet/Keystore.java |  8 +++-
 .../sparrowwallet/drongo/wallet/Wallet.java   | 45 +++++++++++++++++--
 2 files changed, 48 insertions(+), 5 deletions(-)

diff --git a/src/main/java/com/sparrowwallet/drongo/wallet/Keystore.java b/src/main/java/com/sparrowwallet/drongo/wallet/Keystore.java
index d376359..67cdf5a 100644
--- a/src/main/java/com/sparrowwallet/drongo/wallet/Keystore.java
+++ b/src/main/java/com/sparrowwallet/drongo/wallet/Keystore.java
@@ -104,9 +104,13 @@ public class Keystore {
         return new ExtendedKey(derivedKey, derivedKey.getParentFingerprint(), derivation.get(derivation.size() - 1));
     }
 
+    public DeterministicKey getKey(WalletNode walletNode) {
+        return getKey(walletNode.getKeyPurpose(), walletNode.getIndex());
+    }
+
     public DeterministicKey getKey(KeyPurpose keyPurpose, int keyIndex) {
-        List<ChildNumber> receivingDerivation = List.of(extendedPublicKey.getKeyChildNumber(), keyPurpose.getPathIndex(), new ChildNumber(keyIndex));
-        return extendedPublicKey.getKey(receivingDerivation);
+        List<ChildNumber> derivation = List.of(extendedPublicKey.getKeyChildNumber(), keyPurpose.getPathIndex(), new ChildNumber(keyIndex));
+        return extendedPublicKey.getKey(derivation);
     }
 
     public KeyDerivation getDerivation(KeyPurpose keyPurpose, int keyIndex) {
diff --git a/src/main/java/com/sparrowwallet/drongo/wallet/Wallet.java b/src/main/java/com/sparrowwallet/drongo/wallet/Wallet.java
index 11190f7..3490da5 100644
--- a/src/main/java/com/sparrowwallet/drongo/wallet/Wallet.java
+++ b/src/main/java/com/sparrowwallet/drongo/wallet/Wallet.java
@@ -12,6 +12,7 @@ import com.sparrowwallet.drongo.psbt.PSBT;
 import com.sparrowwallet.drongo.psbt.PSBTInput;
 
 import java.util.*;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 
@@ -569,6 +570,19 @@ public class Wallet {
     }
 
     public boolean canSign(PSBT psbt) {
+        return !getSigningNodes(psbt).isEmpty();
+    }
+
+    /**
+     * Determines which nodes in this wallet can sign which inputs in the provided PSBT
+     *
+     * @param psbt The PSBT to be signed
+     * @return A map if the PSBT inputs and the nodes that can sign them
+     */
+    public Map<PSBTInput, WalletNode> getSigningNodes(PSBT psbt) {
+        Map<PSBTInput, WalletNode> signingNodes = new LinkedHashMap<>();
+        Map<Script, WalletNode> walletOutputScripts = getWalletOutputScripts();
+
         for(int inputIndex = 0; inputIndex < psbt.getTransaction().getInputs().size(); inputIndex++) {
             TransactionInput txInput = psbt.getTransaction().getInputs().get(inputIndex);
             PSBTInput psbtInput = psbt.getPsbtInputs().get(inputIndex);
@@ -580,13 +594,38 @@ public class Wallet {
 
             if(utxo != null) {
                 Script scriptPubKey = utxo.getScript();
-                if(isWalletOutputScript(scriptPubKey)) {
-                    return true;
+                WalletNode signingNode = walletOutputScripts.get(scriptPubKey);
+                if(signingNode != null) {
+                    signingNodes.put(psbtInput, signingNode);
                 }
             }
         }
 
-        return false;
+        return signingNodes;
+    }
+
+    /**
+     * Determines which keystores have signed a partially signed (unfinalized) PSBT
+     *
+     * @param psbt The partially signed PSBT
+     */
+    public Map<PSBTInput, List<Keystore>> getSignedKeystores(PSBT psbt) {
+        Map<PSBTInput, WalletNode> signingNodes = getSigningNodes(psbt);
+        Map<PSBTInput, List<Keystore>> signedKeystores = new LinkedHashMap<>();
+
+        for(PSBTInput psbtInput : signingNodes.keySet()) {
+            WalletNode walletNode = signingNodes.get(psbtInput);
+            Map<ECKey, Keystore> keystoreKeysForNode = getKeystores().stream().collect(Collectors.toMap(keystore -> keystore.getKey(walletNode), Function.identity(),
+                    (u, v) -> { throw new IllegalStateException("Duplicate keys from different keystores for node " + walletNode.getDerivationPath()); },
+                    LinkedHashMap::new));
+
+            keystoreKeysForNode.keySet().retainAll(psbtInput.getPartialSignatures().keySet());
+
+            List<Keystore> inputSignedKeystores = new ArrayList<>(keystoreKeysForNode.values());
+            signedKeystores.put(psbtInput, inputSignedKeystores);
+        }
+
+        return signedKeystores;
     }
 
     public BitcoinUnit getAutoUnit() {