deduplicate signed keystores on signature

src/main/java/com/sparrowwallet/drongo/protocol/TransactionSignature.java

@@ -5,6 +5,7 @@ import com.sparrowwallet.drongo.crypto.ECKey;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.math.BigInteger;
+import java.util.Objects;
 
 public class TransactionSignature extends ECKey.ECDSASignature {
     /**
@@ -131,6 +132,26 @@ public class TransactionSignature extends ECKey.ECDSASignature {
         return new TransactionSignature(super.toCanonicalised(), getSigHash());
     }
 
+    @Override
+    public boolean equals(Object o) {
+        if(this == o) {
+            return true;
+        }
+        if(o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        if(!super.equals(o)) {
+            return false;
+        }
+        TransactionSignature signature = (TransactionSignature) o;
+        return sighashFlags == signature.sighashFlags;
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(super.hashCode(), sighashFlags);
+    }
+
     /**
      * Returns a decoded signature.
      *

src/main/java/com/sparrowwallet/drongo/psbt/PSBTInput.java

@@ -452,6 +452,26 @@ public class PSBTInput {
         return false;
     }
 
+    public Map<ECKey, TransactionSignature> getSigningKeys(Set<ECKey> availableKeys) {
+        Collection<TransactionSignature> signatures = getSignatures();
+        Script signingScript = getSigningScript();
+
+        Map<ECKey, TransactionSignature> signingKeys = new LinkedHashMap<>();
+        if(signingScript != null) {
+            Sha256Hash hash = getHashForSignature(signingScript, getSigHash() == null ? SigHash.ALL : getSigHash());
+
+            for(ECKey sigPublicKey : availableKeys) {
+                for(TransactionSignature signature : signatures) {
+                    if(sigPublicKey.verify(hash, signature)) {
+                        signingKeys.put(sigPublicKey, signature);
+                    }
+                }
+            }
+        }
+
+        return signingKeys;
+    }
+
     public ScriptType getScriptType() {
         Script signingScript = getUtxo().getScript();
 

src/main/java/com/sparrowwallet/drongo/wallet/FinalizingPSBTWallet.java

@@ -44,7 +44,7 @@ public class FinalizingPSBTWallet extends Wallet {
             signedNodeKeys.put(signedNode, new ArrayList<>(keys));
 
             ScriptType scriptType = psbtInput.getScriptType();
-            if(scriptType == null || (getScriptType() != null && scriptType.equals(getScriptType()))) {
+            if(scriptType == null || (getScriptType() != null && !scriptType.equals(getScriptType()))) {
                 throw new IllegalArgumentException("Cannot determine the script type from the PSBT, or there are multiple script types");
             } else {
                 setScriptType(scriptType);
@@ -71,12 +71,17 @@ public class FinalizingPSBTWallet extends Wallet {
     }
 
     @Override
-    public Map<PSBTInput, List<Keystore>> getSignedKeystores(PSBT psbt) {
+    public Map<PSBTInput, Map<TransactionSignature, Keystore>> getSignedKeystores(PSBT psbt) {
-        Map<PSBTInput, List<Keystore>> signedKeystores = new LinkedHashMap<>();
+        Map<PSBTInput, Map<TransactionSignature, Keystore>> signedKeystores = new LinkedHashMap<>();
         for(PSBTInput psbtInput : psbt.getPsbtInputs()) {
-            Collection<TransactionSignature> signatures = psbtInput.getSignatures();
+            List<TransactionSignature> signatures = new ArrayList<>(psbtInput.getSignatures());
-            signedKeystores.put(psbtInput, IntStream.range(1, signatures.size() + 1).mapToObj(i -> new Keystore("Keystore " + i)).collect(Collectors.toList()));
+            Map<TransactionSignature, Keystore> signatureKeystoreMap = new LinkedHashMap<>();
+            for(int i = 0; i < signatures.size(); i++) {
+                signatureKeystoreMap.put(signatures.get(i), new Keystore("Keystore " + (i + 1)));
             }
+            signedKeystores.put(psbtInput, signatureKeystoreMap);
+        }
+
         return signedKeystores;
     }
 

src/main/java/com/sparrowwallet/drongo/wallet/Wallet.java

@@ -611,13 +611,14 @@ public class Wallet {
     }
 
     /**
-     * Determines which keystores have signed a partially signed (unfinalized) PSBT
+     * Determines which keystores have signed a PSBT
      *
-     * @param psbt The partially signed PSBT
+     * @param psbt The partially signed or finalized PSBT
+     * @return A map keyed with the PSBTInput mapped to a map of the signatures and associated keystores that signed it
      */
-    public Map<PSBTInput, List<Keystore>> getSignedKeystores(PSBT psbt) {
+    public Map<PSBTInput, Map<TransactionSignature, Keystore>> getSignedKeystores(PSBT psbt) {
         Map<PSBTInput, WalletNode> signingNodes = getSigningNodes(psbt);
-        Map<PSBTInput, List<Keystore>> signedKeystores = new LinkedHashMap<>();
+        Map<PSBTInput, Map<TransactionSignature, Keystore>> signedKeystores = new LinkedHashMap<>();
 
         for(PSBTInput psbtInput : signingNodes.keySet()) {
             WalletNode walletNode = signingNodes.get(psbtInput);
@@ -625,10 +626,21 @@ public class Wallet {
                     (u, v) -> { throw new IllegalStateException("Duplicate keys from different keystores for node " + walletNode.getDerivationPath()); },
                     LinkedHashMap::new));
 
-            keystoreKeysForNode.keySet().retainAll(psbtInput.getPartialSignatures().keySet());
+            Map<ECKey, TransactionSignature> keySignatureMap;
+            if(psbt.isFinalized()) {
+                keySignatureMap = psbtInput.getSigningKeys(keystoreKeysForNode.keySet());
+            } else {
+                keySignatureMap = psbtInput.getPartialSignatures();
+            }
 
-            List<Keystore> inputSignedKeystores = new ArrayList<>(keystoreKeysForNode.values());
+            keystoreKeysForNode.keySet().retainAll(keySignatureMap.keySet());
-            signedKeystores.put(psbtInput, inputSignedKeystores);
+
+            Map<TransactionSignature, Keystore> inputSignatureKeystores = new LinkedHashMap<>();
+            for(ECKey signingKey : keystoreKeysForNode.keySet()) {
+                inputSignatureKeystores.put(keySignatureMap.get(signingKey), keystoreKeysForNode.get(signingKey));
+            }
+
+            signedKeystores.put(psbtInput, inputSignatureKeystores);
         }
 
         return signedKeystores;