psbt and scripting improvments

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

@@ -142,7 +142,7 @@ public class Script {
      */
     public Address[] getToAddresses() throws NonStandardScriptException {
         if (ScriptPattern.isP2PK(this))
-            return new Address[] { new P2PKAddress( ScriptPattern.extractPKFromP2PK(this)) };
+            return new Address[] { new P2PKAddress( ScriptPattern.extractPKFromP2PK(this).getPubKey()) };
         else if (ScriptPattern.isP2PKH(this))
             return new Address[] { new P2PKHAddress( ScriptPattern.extractHashFromP2PKH(this)) };
         else if (ScriptPattern.isP2SH(this))
@@ -179,6 +179,17 @@ public class Script {
         return null;
     }
 
+    public List<TransactionSignature> getSignatures() {
+        List<TransactionSignature> signatures = new ArrayList<>();
+        for(ScriptChunk chunk : chunks) {
+            if(chunk.isSignature()) {
+                signatures.add(chunk.getSignature());
+            }
+        }
+
+        return signatures;
+    }
+
     public static int decodeFromOpN(int opcode) {
         if((opcode != OP_0 && opcode != OP_1NEGATE) && (opcode < OP_1 || opcode > OP_16)) {
             throw new ProtocolException("decodeFromOpN called on non OP_N opcode: " + opcode);
@@ -271,7 +282,7 @@ public class Script {
             if(chunk.isSignature()) {
                 builder.append("<signature").append(signatureCount++).append(">");
             } else if(chunk.isScript()) {
-                Script nestedScript = new Script(chunk.getData());
+                Script nestedScript = chunk.getScript();
                 if(ScriptPattern.isP2WPKH(nestedScript)) {
                     builder.append("(OP_0 <wpkh>)");
                 } else if(ScriptPattern.isP2WSH(nestedScript)) {

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

@@ -96,6 +96,14 @@ public class ScriptChunk {
         return true;
     }
 
+    public TransactionSignature getSignature() {
+        try {
+            return TransactionSignature.decodeFromBitcoin(data, false, false);
+        } catch(SignatureDecodeException e) {
+            throw new ProtocolException("Could not decode signature", e);
+        }
+    }
+
     public boolean isScript() {
         if(data == null || data.length == 0) {
             return false;
@@ -110,6 +118,10 @@ public class ScriptChunk {
         return true;
     }
 
+    public Script getScript() {
+        return new Script(data);
+    }
+
     public boolean isPubKey() {
         if(data == null || data.length == 0) {
             return false;
@@ -118,6 +130,10 @@ public class ScriptChunk {
         return ECKey.isPubKey(data);
     }
 
+    public ECKey getPubKey() {
+        return ECKey.fromPublicOnly(data);
+    }
+
     public byte[] toByteArray() {
         ByteArrayOutputStream stream = new ByteArrayOutputStream();
         try {

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

@@ -2,6 +2,7 @@ package com.sparrowwallet.drongo.protocol;
 
 import com.sparrowwallet.drongo.address.Address;
 import com.sparrowwallet.drongo.address.P2PKAddress;
+import com.sparrowwallet.drongo.crypto.ECKey;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -34,8 +35,8 @@ public class ScriptPattern {
      * Extract the pubkey from a P2PK scriptPubKey. It's important that the script is in the correct form, so you
      * will want to guard calls to this method with {@link #isP2PK(Script)}.
      */
-    public static byte[] extractPKFromP2PK(Script script) {
-        return script.chunks.get(0).data;
+    public static ECKey extractPKFromP2PK(Script script) {
+        return ECKey.fromPublicOnly(script.chunks.get(0).data);
     }
 
     /**
@@ -106,6 +107,14 @@ public class ScriptPattern {
         return true;
     }
 
+    /**
+     * Extract the script hash from a P2SH scriptPubKey. It's important that the script is in the correct form, so you
+     * will want to guard calls to this method with {@link #isP2SH(Script)}.
+     */
+    public static byte[] extractHashFromP2SH(Script script) {
+        return script.chunks.get(1).data;
+    }
+
     /**
      * Returns whether this script matches the format used for multisig outputs:
      * {@code [n] [keys...] [m] CHECKMULTISIG}
@@ -150,14 +159,6 @@ public class ScriptPattern {
         return addresses.toArray(new Address[addresses.size()]);
     }
 
-    /**
-     * Extract the script hash from a P2SH scriptPubKey. It's important that the script is in the correct form, so you
-     * will want to guard calls to this method with {@link #isP2SH(Script)}.
-     */
-    public static byte[] extractHashFromP2SH(Script script) {
-        return script.chunks.get(1).data;
-    }
-
     /**
      * Returns true if this script is of the form {@code OP_0 <hash[20]>}. This is a P2WPKH scriptPubKey.
      */

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

@@ -165,9 +165,11 @@ public class Transaction extends TransactionPart {
         // script_witnesses
         if (useSegwit) {
             for (TransactionInput in : inputs) {
+                if (in.hasWitness()) {
                     in.getWitness().bitcoinSerializeToStream(stream);
                 }
             }
+        }
         // lock_time
         uint32ToByteStreamLE(locktime, stream);
     }
@@ -226,14 +228,9 @@ public class Transaction extends TransactionPart {
     private void parseWitnesses() {
         int numWitnesses = inputs.size();
         for (int i = 0; i < numWitnesses; i++) {
-            long pushCount = readVarInt();
-            TransactionWitness witness = new TransactionWitness((int) pushCount);
+            TransactionWitness witness = new TransactionWitness(this, rawtx, cursor);
             inputs.get(i).setWitness(witness);
-            for (int y = 0; y < pushCount; y++) {
-                long pushSize = readVarInt();
-                byte[] push = readBytes((int) pushSize);
-                witness.setPush(y, push);
-            }
+            cursor += witness.getLength();
         }
     }
 
@@ -261,9 +258,11 @@ public class Transaction extends TransactionPart {
         // script_witnesses
         if(isSegwit()) {
             for (TransactionInput in : inputs) {
+                if (in.hasWitness()) {
                     wu += in.getWitness().getLength();
                 }
             }
+        }
         // lock_time
         wu += 4*4;
 

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

@@ -66,7 +66,7 @@ public class TransactionInput extends TransactionPart {
     }
 
     public TransactionWitness getWitness() {
-        return witness != null ? witness : TransactionWitness.EMPTY;
+        return witness;
     }
 
     public void setWitness(TransactionWitness witness) {

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

@@ -1,8 +1,8 @@
 package com.sparrowwallet.drongo.protocol;
 
-import com.sparrowwallet.drongo.Utils;
 import org.bouncycastle.util.encoders.Hex;
 
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.OutputStream;
 import java.util.ArrayList;
@@ -10,20 +10,35 @@ import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 
-public class TransactionWitness {
-    public static final TransactionWitness EMPTY = new TransactionWitness(0);
+public class TransactionWitness extends TransactionPart {
+    private List<byte[]> pushes;
 
-    private final List<byte[]> pushes;
+    public TransactionWitness(Transaction parent, byte[] rawtx, int offset) {
+        super(rawtx, offset);
+        setParent(parent);
+        if(pushes == null) {
+            pushes = new ArrayList<>();
+        }
+    }
 
-    public TransactionWitness(int pushCount) {
-        pushes = new ArrayList<>(Math.min(pushCount, Utils.MAX_INITIAL_ARRAY_LENGTH));
+    protected void parse() throws ProtocolException {
+        long pushCount = readVarInt();
+        for (int y = 0; y < pushCount; y++) {
+            long pushSize = readVarInt();
+            byte[] push = readBytes((int)pushSize);
+            setPush(y, push);
+        }
     }
 
     public List<byte[]> getPushes() {
         return Collections.unmodifiableList(pushes);
     }
 
-    public void setPush(int i, byte[] value) {
+    protected void setPush(int i, byte[] value) {
+        if(pushes == null) {
+            pushes = new ArrayList<>();
+        }
+
         while (i >= pushes.size()) {
             pushes.add(new byte[]{});
         }
@@ -54,6 +69,15 @@ public class TransactionWitness {
         }
     }
 
+    public byte[] toByteArray() {
+        ByteArrayOutputStream baos = new ByteArrayOutputStream();
+        try {
+            bitcoinSerializeToStream(baos);
+        } catch(IOException e) { }
+
+        return baos.toByteArray();
+    }
+
     @Override
     public String toString() {
         StringBuilder builder = new StringBuilder();
@@ -80,6 +104,27 @@ public class TransactionWitness {
         return scriptChunks;
     }
 
+    public List<TransactionSignature> getSignatures() {
+        List<TransactionSignature> signatures = new ArrayList<>();
+        List<ScriptChunk> scriptChunks = this.asScriptChunks();
+        for(ScriptChunk chunk : scriptChunks) {
+            if(chunk.isSignature()) {
+                signatures.add(chunk.getSignature());
+            }
+        }
+
+        return signatures;
+    }
+
+    public Script getWitnessScript() {
+        List<ScriptChunk> scriptChunks = this.asScriptChunks();
+        if(scriptChunks.get(scriptChunks.size() - 1).isScript()) {
+            return scriptChunks.get(scriptChunks.size() - 1).getScript();
+        }
+
+        return null;
+    }
+
     @Override
     public boolean equals(Object o) {
         if (this == o) return true;

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

@@ -11,10 +11,7 @@ import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.nio.charset.StandardCharsets;
-import java.util.Arrays;
-import java.util.LinkedHashMap;
-import java.util.List;
-import java.util.Map;
+import java.util.*;
 
 import static com.sparrowwallet.drongo.psbt.PSBTEntry.parseKeyDerivation;
 
@@ -39,7 +36,7 @@ public class PSBTInput {
     private Script witnessScript;
     private Map<ECKey, KeyDerivation> derivedPublicKeys = new LinkedHashMap<>();
     private Script finalScriptSig;
-    private Script finalScriptWitness;
+    private TransactionWitness finalScriptWitness;
     private String porCommitment;
     private Map<String, String> proprietary = new LinkedHashMap<>();
 
@@ -162,9 +159,9 @@ public class PSBTInput {
                     break;
                 case PSBT_IN_FINAL_SCRIPTWITNESS:
                     entry.checkOneByteKey();
-                    Script finalScriptWitness = new Script(entry.getData());
+                    TransactionWitness finalScriptWitness = new TransactionWitness(null, entry.getData(), 0);
                     this.finalScriptWitness = finalScriptWitness;
-                    log.debug("Found input final scriptWitness script hex " + Hex.toHexString(finalScriptWitness.getProgram()) + " script " + finalScriptWitness.toString());
+                    log.debug("Found input final scriptWitness " + finalScriptWitness.toString());
                     break;
                 case PSBT_IN_POR_COMMITMENT:
                     entry.checkOneByteKey();
@@ -217,7 +214,7 @@ public class PSBTInput {
         return finalScriptSig;
     }
 
-    public Script getFinalScriptWitness() {
+    public TransactionWitness getFinalScriptWitness() {
         return finalScriptWitness;
     }
 
@@ -229,6 +226,18 @@ public class PSBTInput {
         return partialSignatures;
     }
 
+    public ECKey getKeyForSignature(TransactionSignature signature) {
+        if(partialSignatures != null) {
+            for(Map.Entry<ECKey, TransactionSignature> entry : partialSignatures.entrySet()) {
+                if(entry.getValue().equals(signature)) {
+                    return entry.getKey();
+                }
+            }
+        }
+
+        return null;
+    }
+
     public Map<ECKey, KeyDerivation> getDerivedPublicKeys() {
         return derivedPublicKeys;
     }
@@ -263,6 +272,8 @@ public class PSBTInput {
                     }
                 }
 
+                //TODO: Implement Bitcoin Script engine to verify finalScriptSig and finalScriptWitness
+
                 return true;
             }
         }
@@ -275,10 +286,12 @@ public class PSBTInput {
         Script signingScript = getNonWitnessUtxo() != null ? getNonWitnessUtxo().getOutputs().get(vout).getScript() : getWitnessUtxo().getScript();
 
         if(ScriptPattern.isP2SH(signingScript)) {
-            if(getRedeemScript() == null) {
-                return null;
-            } else {
+            if(getRedeemScript() != null) {
                 signingScript = getRedeemScript();
+            } else if(getFinalScriptSig() != null) {
+                signingScript = getFinalScriptSig().getFirstNestedScript();
+            } else {
+                return null;
             }
         }
 
@@ -286,10 +299,15 @@ public class PSBTInput {
             Address address = new P2PKHAddress(signingScript.getPubKeyHash());
             signingScript = address.getOutputScript();
         } else if(ScriptPattern.isP2WSH(signingScript)) {
-            if(getWitnessScript() == null) {
-                return null;
-            } else {
+            if(getWitnessScript() != null) {
                 signingScript = getWitnessScript();
+            } else if(getFinalScriptWitness() != null) {
+                List<ScriptChunk> witnessChunks = getFinalScriptWitness().asScriptChunks();
+                if(witnessChunks.get(witnessChunks.size() - 1).isScript()) {
+                    return witnessChunks.get(witnessChunks.size() - 1).getScript();
+                }
+            } else {
+                return null;
             }
         }
 

src/main/java/module-info.java

@@ -5,4 +5,5 @@ module com.sparrowwallet.drongo {
     exports com.sparrowwallet.drongo.psbt;
     exports com.sparrowwallet.drongo.protocol;
     exports com.sparrowwallet.drongo.address;
+    exports com.sparrowwallet.drongo.crypto;
 }

src/test/java/com/sparrowwallet/drongo/psbt/PSBTTest.java

@@ -284,6 +284,6 @@ public class PSBTTest {
 
         Assert.assertEquals("00473044022074018ad4180097b873323c0015720b3684cc8123891048e7dbcd9b55ad679c99022073d369b740e3eb53dcefa33823c8070514ca55a7dd9544f157c167913261118c01483045022100f61038b308dc1da865a34852746f015772934208c6d24454393cd99bdf2217770220056e675a675a6d0a02b85b14e5e29074d8a25a9b5760bea2816f661910a006ea01475221029583bf39ae0a609747ad199addd634fa6108559d6c5cd39b4c2183f1ab96e07f2102dab61ff49a14db6a7d02b0cd1fbb78fc4b18312b5b4e54dae4dba2fbfef536d752ae", psbt1.getPsbtInputs().get(0).getFinalScriptSig().getProgramAsHex());
         Assert.assertEquals("2200208c2353173743b595dfb4a07b72ba8e42e3797da74e87fe7d9d7497e3b2028903", psbt1.getPsbtInputs().get(1).getFinalScriptSig().getProgramAsHex());
-        Assert.assertEquals("0400473044022062eb7a556107a7c73f45ac4ab5a1dddf6f7075fb1275969a7f383efff784bcb202200c05dbb7470dbf2f08557dd356c7325c1ed30913e996cd3840945db12228da5f01473044022065f45ba5998b59a27ffe1a7bed016af1f1f90d54b3aa8f7450aa5f56a25103bd02207f724703ad1edb96680b284b56d4ffcb88f7fb759eabbe08aa30f29b851383d20147522103089dc10c7ac6db54f91329af617333db388cead0c231f723379d1b99030b02dc21023add904f3d6dcf59ddb906b0dee23529b7ffb9ed50e5e86151926860221f0e7352ae", psbt1.getPsbtInputs().get(1).getFinalScriptWitness().getProgramAsHex());
+        Assert.assertEquals("0400473044022062eb7a556107a7c73f45ac4ab5a1dddf6f7075fb1275969a7f383efff784bcb202200c05dbb7470dbf2f08557dd356c7325c1ed30913e996cd3840945db12228da5f01473044022065f45ba5998b59a27ffe1a7bed016af1f1f90d54b3aa8f7450aa5f56a25103bd02207f724703ad1edb96680b284b56d4ffcb88f7fb759eabbe08aa30f29b851383d20147522103089dc10c7ac6db54f91329af617333db388cead0c231f723379d1b99030b02dc21023add904f3d6dcf59ddb906b0dee23529b7ffb9ed50e5e86151926860221f0e7352ae", psbt1.getPsbtInputs().get(1).getFinalScriptWitness().toByteArray());
     }
 }