From d6f45c351c875cc6c8b07016caf9953645b307d5 Mon Sep 17 00:00:00 2001
From: davidchocholaty <david.chocholaty@seznam.cz>
Date: Wed, 6 Nov 2024 23:05:47 +0100
Subject: [PATCH] Add p2wpkh

---
 src/script.py      |  9 +++---
 src/transaction.py | 69 ++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 71 insertions(+), 7 deletions(-)

diff --git a/src/script.py b/src/script.py
index 8c34db3..b62cc79 100644
--- a/src/script.py
+++ b/src/script.py
@@ -35,16 +35,17 @@ def is_empty(self) -> bool:
         return len(self._items) == 0
 
 class Script:
-    def __init__(self, script: bytes, json_transaction: dict = None, input_index: int = 0):
+    def __init__(self, script: bytes, json_transaction: dict = None, input_index: int = 0, segwit: bool = False):
         self.script = script
         self.stack = Stack()
         self.alt_stack = Stack()
         self.if_stack: List[bool] = []
         self.transaction = json_transaction  # Store JSON transaction
         self.input_index = input_index
+        self.segwit = segwit
         
     def create_signature_hash(self, hash_type: int) -> bytes:
-        data_signed = serialize_transaction(self.transaction, self.input_index, int(hash_type))
+        data_signed = serialize_transaction(self.transaction, self.input_index, int(hash_type), self.segwit)
         return hashlib.sha256(data_signed).digest()
         """
         Create the signature hash for the transaction based on the hash type.
@@ -179,7 +180,7 @@ def execute(self) -> bool:
                 i += self._execute_opcode(op_name)
             
             # Script executed successfully if stack is not empty and top value is true
-            if self.stack.is_empty():
+            if self.stack.is_empty():                
                 return False
             
             return self.stack.pop() != b'\x00'
@@ -191,7 +192,7 @@ def _execute_opcode(self, op_name: str) -> int:
         """Execute a single opcode and return how many bytes to advance"""
         
         # Constants
-        if op_name == 'OP_0':            
+        if op_name == 'OP_0':
             self.stack.push(b'\x00')
             return 1
         elif op_name == 'OP_1NEGATE':
diff --git a/src/transaction.py b/src/transaction.py
index 45d13ba..661fedd 100644
--- a/src/transaction.py
+++ b/src/transaction.py
@@ -110,7 +110,7 @@ def valid_input(self, vin_idx, vin):
         elif scriptpubkey_type == "v1_p2tr":
             pass
         elif scriptpubkey_type == "v0_p2wpkh":
-            return self.validate_p2wpkh(vin)
+            return self.validate_p2wpkh(vin_idx, vin)
         
         # Unknown script type.
         return False        
@@ -139,5 +139,68 @@ def validate_p2pkh(self, vin_idx, vin):
 
         return is_valid
 
-    def validate_p2wpkh(self, vin):
-        return False
+    def validate_p2wpkh(self, vin_idx, vin):
+        """
+        Validate a Pay-to-Witness-Public-Key-Hash (P2WPKH) transaction input
+        
+        Args:
+            vin_idx (int): Index of the input being validated
+            vin (dict): Input data containing witness information
+            
+        Returns:
+            bool: True if the P2WPKH input is valid, False otherwise
+        """
+        # Check for witness data
+        witness = vin.get("witness", [])
+        if len(witness) != 2:  # P2WPKH requires exactly 2 witness items (signature and pubkey)
+            return False
+            
+        # Get witness components
+        signature = bytes.fromhex(witness[0])
+        pubkey = bytes.fromhex(witness[1])
+        
+        # Get previous output's scriptPubKey
+        prevout = vin.get("prevout", {})
+        if not prevout:
+            return False
+            
+        scriptpubkey = bytes.fromhex(prevout.get("scriptpubkey", ""))
+        
+        # Verify the script is proper P2WPKH format (OP_0 <20-byte-key-hash>)
+        if len(scriptpubkey) != 22 or scriptpubkey[0] != 0x00 or scriptpubkey[1] != 0x14:
+            return False
+        
+        # Create P2PKH-style script from witness data
+        # P2WPKH witness program is executed as if it were P2PKH scriptPubKey
+        p2pkh_script = (
+            bytes([len(signature)]) +  # Push signature
+            signature +
+            bytes([len(pubkey)]) +     # Push pubkey
+            pubkey +
+            # Standard P2PKH script operations
+            bytes([
+                0x76,  # OP_DUP
+                0xa9,  # OP_HASH160
+                0x14   # Push 20 bytes
+            ]) +
+            scriptpubkey[2:22] +  # 20-byte pubkey hash from witness program
+            bytes([
+                0x88,  # OP_EQUALVERIFY
+                0xac   # OP_CHECKSIG
+            ])
+        )
+        
+        # Create script object with witness-specific transaction data
+        script = Script(
+            p2pkh_script,
+            json_transaction=self.json_transaction,
+            input_index=vin_idx,
+            segwit=True
+        )
+        
+        # Execute the script
+        try:
+            return script.execute()
+        except Exception as e:
+            print(f"P2WPKH validation error: {str(e)}")
+            return False