p2tr_util.py

import hashlib
import os
import sys
from typing import Optional, Union, List, Tuple
from enum import Enum
import secp256k1 as ice

FieldOrder = 115792089237316195423570985008687907853269984665640564039457584007908834671663
GroupOrder = 115792089237316195423570985008687907852837564279074904382605163141518161494337

def get_x(point):
    point = point.hex()
    return int(point[2:66], 16)

def get_y(point):
    point = point.hex()
    return int(point[66:], 16)
    
class Encoding(Enum):
    """Enumeration type to list the various supported encodings."""
    BECH32 = 1
    BECH32M = 2


CHARSET = "qpzry9x8gf2tvdw0s3jn54khce6mua7l"
BECH32M_CONST = 0x2bc830a3


def bech32_polymod(values: List[int]) -> int:
    """Internal function that computes the Bech32 checksum."""
    generator = [0x3b6a57b2, 0x26508e6d, 0x1ea119fa, 0x3d4233dd, 0x2a1462b3]
    chk = 1
    for value in values:
        top = chk >> 25
        chk = (chk & 0x1ffffff) << 5 ^ value
        for i in range(5):
            chk ^= generator[i] if ((top >> i) & 1) else 0
    return chk


def bech32_hrp_expand(hrp: str) -> List[int]:
    """Expand the HRP into values for checksum computation."""
    return [ord(x) >> 5 for x in hrp] + [0] + [ord(x) & 31 for x in hrp]


def bech32_verify_checksum(hrp: str, data: List[int]) -> Union[Encoding, None]:
    """Verify a checksum given HRP and converted data characters."""
    const = bech32_polymod(bech32_hrp_expand(hrp) + data)
    if const == 1:
        return Encoding.BECH32
    if const == BECH32M_CONST:
        return Encoding.BECH32M
    return None


def bech32_create_checksum(hrp: str, data: List[int], spec: Encoding) -> List[int]:
    """Compute the checksum values given HRP and data."""
    values = bech32_hrp_expand(hrp) + data
    const = BECH32M_CONST if spec == Encoding.BECH32M else 1
    polymod = bech32_polymod(values + [0, 0, 0, 0, 0, 0]) ^ const
    return [(polymod >> 5 * (5 - i)) & 31 for i in range(6)]


def bech32_encode(hrp: str, data: List[int], spec: Encoding) -> str:
    """Compute a Bech32 string given HRP and data values."""
    combined = data + bech32_create_checksum(hrp, data, spec)
    return hrp + '1' + ''.join([CHARSET[d] for d in combined])


def bech32_decode(bech: str) -> Union[Tuple[None, None, None], Tuple[str, List[int], Encoding]]:
    """Validate a Bech32/Bech32m string, and determine HRP and data."""
    if ((any(ord(x) < 33 or ord(x) > 126 for x in bech)) or
            (bech.lower() != bech and bech.upper() != bech)):
        return None, None, None
    bech = bech.lower()
    pos = bech.rfind('1')
    if pos < 1 or pos + 7 > len(bech) or len(bech) > 90:
        return None, None, None
    if not all(x in CHARSET for x in bech[pos+1:]):
        return None, None, None
    hrp = bech[:pos]
    data = [CHARSET.find(x) for x in bech[pos+1:]]
    spec = bech32_verify_checksum(hrp, data)
    if spec is None:
        return None, None, None
    return hrp, data[:-6], spec


def convertbits(data: bytes, frombits: int, tobits: int, pad: bool = True) -> Optional[List[int]]:
    """General power-of-2 base conversion."""
    acc = 0
    bits = 0
    ret = []
    maxv = (1 << tobits) - 1
    max_acc = (1 << (frombits + tobits - 1)) - 1
    for value in data:
        if value < 0 or (value >> frombits):
            return None
        acc = ((acc << frombits) | value) & max_acc
        bits += frombits
        while bits >= tobits:
            bits -= tobits
            ret.append((acc >> bits) & maxv)
    if pad:
        if bits:
            ret.append((acc << (tobits - bits)) & maxv)
    elif bits >= frombits or ((acc << (tobits - bits)) & maxv):
        return None
    return ret


def decode(hrp: str, addr: str) -> Union[Tuple[None, None], Tuple[int, List[int]]]:
    """Decode a segwit address."""
    hrpgot, data, spec = bech32_decode(addr)
    if hrpgot != hrp:
        return None, None
    decoded = convertbits(data[1:], 5, 8, False)
    # Witness programs are between 2 and 40 bytes in length.
    if decoded is None or len(decoded) < 2 or len(decoded) > 40:
        return None, None
    # Witness versions are in range 0..16.
    if data[0] > 16:
        return None, None
    # Witness v0 programs must be exactly length 20 or 32.
    if data[0] == 0 and len(decoded) != 20 and len(decoded) != 32:
        return None, None
    # Witness v0 uses Bech32; v1 through v16 use Bech32m.
    if data[0] == 0 and spec != Encoding.BECH32 or data[0] != 0 and spec != Encoding.BECH32M:
        return None, None
    # Success.
    return data[0], decoded


def encode(hrp: str, witver: int, witprog: bytes) -> Optional[str]:
    """Encode a segwit address."""
    spec = Encoding.BECH32 if witver == 0 else Encoding.BECH32M
    ret = bech32_encode(hrp, [witver] + convertbits(witprog, 8, 5), spec)
    if decode(hrp, ret) == (None, None):
        return None
    return ret


def pubkey_to_segwit_v1_addr(human_readable_part: str, witprog: bytes) -> Optional[str]:
    """ Derives bech32m (p2tr) address from pubkey """
    witver = 1
    addr = encode(human_readable_part, witver, witprog)
    return addr

def bytes_from_int(x: int) -> bytes:
    return x.to_bytes(32, byteorder="big")


def int_from_bytes(b: bytes) -> int:
    return int.from_bytes(b, byteorder="big")


def tagged_hash(tag: str, msg: bytes) -> bytes:
    tag_hash = hashlib.sha256(tag.encode()).digest()
    return hashlib.sha256(tag_hash + tag_hash + msg).digest()


def lift_x(x: bytes):
    x = int_from_bytes(x)
    y_sq = (pow(x, 3, FieldOrder) + 7) % FieldOrder
    y1 = pow(y_sq, (FieldOrder + 1) // 4, FieldOrder)
    y2 = FieldOrder - y1
    if y1 % 2 == 0:
        y_coord = y1
    else:
        y_coord = y2
    P = bytes.fromhex('04'+ hex(x).lstrip('0x').zfill(64)+hex(y_coord).lstrip('0x').zfill(64))
    return P


def has_even_y(point):
    return get_y(point) % 2 == 0


def taproot_tweak_seckey(seckey0: bytes) -> bytes:
    P = ice.scalar_multiplication(int_from_bytes(seckey0))
    seckey = int_from_bytes(seckey0) if has_even_y(P) else GroupOrder - int_from_bytes(seckey0)
    t = int_from_bytes(tagged_hash("TapTweak", bytes_from_int(get_x(P))))
    return bytes_from_int(int((seckey + t) % GroupOrder))


def public_key_x_coordinate_to_taproot_tweaked_pubkey(public_key_x_coordinate: bytes) -> bytes:
    P = lift_x(public_key_x_coordinate)
    t = int_from_bytes(tagged_hash("TapTweak", bytes_from_int(int(public_key_x_coordinate.hex(), 16))))
    Q = ice.point_addition(P,ice.scalar_multiplication(t))
    return bytes_from_int(get_x(Q))