Updating Signing Code for DeSo Balance Model Fork

sign.py

# This file has code for signing a transaction --> Credit: MiniGlome
import hashlib
import hmac
from typing import Union
from hashlib import sha256
from secp256k1 import PrivateKey, PublicKey
import struct

class HmacDrbg:
    def __init__(self, hash_len, q_byte_len):
        if not isinstance(hash_len, int) or hash_len < 2:
            raise ValueError('hash_len must be a positive integer')
        if not isinstance(q_byte_len, int) or q_byte_len < 2:
            raise ValueError('q_byte_len must be a positive integer')
        self.hash_len = hash_len
        self.q_byte_len = q_byte_len
        self.v = bytes([1] * hash_len)
        self.k = bytes([0] * hash_len)
        self.counter = 0
    
    def hmac(self, *values):
        return hmac.digest(self.k, b''.join(values), 'sha256')
    
    def check_sync(self):
        pass # HMAC-SHA256 in Python is already synchronous
    
    def incr(self):
        if self.counter >= 1000:
            raise ValueError('Tried 1,000 k values for sign(), all were invalid')
        self.counter += 1
    
    def reseed(self, seed=b''):
        self.k = self.hmac(self.v, bytes([0x00]), seed)
        self.v = self.hmac(self.v)
        if len(seed) == 0:
            return
        self.k = self.hmac(self.v, bytes([0x01]), seed)
        self.v = self.hmac(self.v)
    
    def reseed_sync(self, seed=b''):
        self.check_sync()
        self.k = self.hmac(self.v, bytes([0x00]), seed)
        self.v = self.hmac(self.v)
        if len(seed) == 0:
            return
        self.k = self.hmac(self.v, bytes([0x01]), seed)
        self.v = self.hmac(self.v)
    
    def generate(self):
        self.incr()
        len_ = 0
        out = []
        while len_ < self.q_byte_len:
            self.v = self.hmac(self.v)
            sl = self.v
            out.append(sl)
            len_ += len(sl)
        return b''.join(out)
    
    def generate_sync(self):
        self.check_sync()
        self.incr()
        len_ = 0
        out = []
        while len_ < self.q_byte_len:
            self.v = self.hmac(self.v)
            sl = self.v
            out.append(sl)
            len_ += len(sl)
        return b''.join(out)


def uvarint64_to_buf(number):
    result = []
    while number >= 0x80:
        result.append((number & 0xff) | 0x80)
        number >>= 7
    result.append(number)
    return bytes(result)

def sign(msg_hash_hex: str, private_key: str):
    msg_hash = bytes.fromhex(msg_hash_hex)
    private_key = PrivateKey(private_key)
    # Prepare signature options
    sig_opts = {
        "canonical": True,
        "der": True,
        "extraEntropy": True,
        "recovered": True,
    }

    # Initialize signature arguments
    seed, m, d = private_key.ecdsa_sign_recoverable(msg_hash)
    sig_args = {"seed": seed, "m": m, "d": d, "extraEntropy": sig_opts["extraEntropy"]}

    # Initialize HmacDrbg
    hash_len = 32  # sha256 hash length
    group_len = 32  # secp256k1 group order length
    drbg = HmacDrbg(hash_len=hash_len, q_byte_len=group_len)

    # Reseed HmacDrbg
    drbg.reseed(seed)

    # Generate signature
    sig = None
    while not sig:
        k = drbg.generate()
        pv = private_key.ecdsa_sign_recoverable(msg_hash, raw=True, k=k)
        sig = private_key.ecdsa_recoverable_serialize(pv)
        sig = HmacDrbg.kmd_to_sig(k, m, d, sig_opts["canonical"], sig)

        # Reseed HmacDrbg if signature is not valid
        if not sig:
            drbg.reseed()

    # Finalize signature
    return private_key.ecdsa_recoverable_convert(sig)

def sha256X2(data):
    if isinstance(data, str):
        data = bytes.fromhex(data)
    sha256_hash = hashlib.sha256(data).digest()
    sha256_hash2 = hashlib.sha256(sha256_hash).digest()
    return sha256_hash2

def inverse_mod(k, p):
    """Returns the inverse of k modulo p.
    This function returns the only integer x such that (x * k) % p == 1.
    k must be non-zero and p must be a prime.
    """
    if k == 0:
        raise ZeroDivisionError("division by zero")

    if k < 0:
        # k ** -1 = p - (-k) ** -1  (mod p)
        return p - inverse_mod(-k, p)

    # Extended Euclidean algorithm.
    s, old_s = 0, 1
    t, old_t = 1, 0
    r, old_r = p, k

    while r != 0:
        quotient = old_r // r
        old_r, r = r, old_r - quotient * r
        old_s, s = s, old_s - quotient * s
        old_t, t = t, old_t - quotient * t

    gcd, x, y = old_r, old_s, old_t

    return x % p


def get_hmac(key, data):
    return hmac.new(key, data, hashlib.sha256).digest()


def hmac_drbg(entropy, string):
    material = entropy + string
    K = b"\x00" * 32
    V = b"\x01" * 32

    K = get_hmac(K, V + b"\x00" + material)
    V = get_hmac(K, V)
    K = get_hmac(K, V + b"\x01" + material)
    V = get_hmac(K, V)

    temp = b""
    while len(temp) < 32:
        V = get_hmac(K, V)
        temp += V

    return temp[:32]


#######


g = (
    0x79BE667EF9DCBBAC55A06295CE870B07029BFCDB2DCE28D959F2815B16F81798,
    0x483ADA7726A3C4655DA4FBFC0E1108A8FD17B448A68554199C47D08FFB10D4B8,
)
p = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEFFFFFC2F
n = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEBAAEDCE6AF48A03BBFD25E8CD0364141


def point_add(point1, point2):
    # Returns the result of point1 + point2 according to the group law.
    if point1 is None:
        return point2
    if point2 is None:
        return point1

    x1, y1 = point1
    x2, y2 = point2

    if x1 == x2 and y1 != y2:
        return None

    if x1 == x2:
        m = (3 * x1 * x1) * inverse_mod(2 * y1, p)
    else:
        m = (y1 - y2) * inverse_mod(x1 - x2, p)

    x3 = m * m - x1 - x2
    y3 = y1 + m * (x3 - x1)
    result = (x3 % p, -y3 % p)

    return result


def scalar_mult(k, point):
    # Returns k * point computed using the double and point_add algorithm.
    result = None
    addend = point

    while k:
        if k & 1:
            # Add.
            result = point_add(result, addend)
        # Double.
        addend = point_add(addend, addend)
        k >>= 1

    return result


#######


def to_DER(r, s):  # Signature to DER format
    r = bytes.fromhex(r)
    s = bytes.fromhex(s)
    if r[0] >= 0x80:
        r = bytes.fromhex("00") + r
    if s[0] >= 0x80:
        s = bytes.fromhex("00") + s
    res = (
        bytes.fromhex("02" + hex(len(r))[2:])
        + r
        + bytes.fromhex("02" + hex(len(s))[2:])
        + s
    )
    res = bytes.fromhex("30" + hex(len(res))[2:]) + res

    return res.hex()


def hexify(n):
    n = hex(n)[2:]
    if len(n) % 2 != 0:
        n = "0" + n
    return n


def Sign_Transaction(seedHex, TransactionHex):
    s256 = hashlib.sha256(
        hashlib.sha256(bytes.fromhex(TransactionHex)).digest()
    ).digest()
    drbg = hmac_drbg(entropy=bytes.fromhex(seedHex), string=s256)
    k = int.from_bytes(drbg, "big")
    kp = scalar_mult(k, g)
    kpX = kp[0]
    r = kpX % n
    s = inverse_mod(k, n) * (r * int(seedHex, 16) + int(s256.hex(), 16))
    s = s % n
    # Enforce low-s
    if s > n // 2:
        s = n - s
    signature = to_DER(hexify(r), hexify(s))
    signed_transaction = (
        TransactionHex[:-2]
        + hex(len(bytearray.fromhex(signature)))[2:]
        + signature
    )

    return signed_transaction

def new_Sign_Transaction(seedHex: str, TransactionHex: str, is_derived_key: bool = False):
    recovery_param = 1
    transactionBytes = bytes.fromhex(TransactionHex)
    print(transactionBytes.hex())
    signatureIndex = len(transactionBytes) - 15 - 2
    print(signatureIndex)
    v1FieldsBuffer = transactionBytes[15:]
    v0FieldsWithoutSignature = transactionBytes[:17]
    print(v0FieldsWithoutSignature.hex())
    hashedTxBytes = sha256X2(transactionBytes)
    print(hashedTxBytes.hex())
    #s256 = hashlib.sha256(
    #    hashlib.sha256(bytes.fromhex(TransactionHex)).digest()
    #).digest()
    #print(s256.hex())
    transactionHashHex = hashedTxBytes.hex()
    privateKey = bytes.fromhex(seedHex)
    print("Right before sign")
    drbg = hmac_drbg(entropy=bytes.fromhex(seedHex), string=hashedTxBytes)
    k = int.from_bytes(drbg, "big")
    kp = scalar_mult(k, g)
    kpX = kp[0]
    r = kpX % n
    s = inverse_mod(k, n) * (r * int(seedHex, 16) + int(hashedTxBytes.hex(), 16))
    s = s % n
    # Enforce low-s
    if s > n // 2:
        s = n - s
    signature = to_DER(hexify(r), hexify(s))
    signed_transaction = (
        TransactionHex[:-2]
        + hex(len(bytearray.fromhex(signature)))[2:]
        + signature
    )
    signatureBytes = bytes.fromhex(signed_transaction)
    #signatureLength = struct.pack("<Q", len("71"))
    signatureLength = uvarint64_to_buf(71)
    print("siglength")
    print(signatureLength)
    if is_derived_key:
        signatureBytes = bytearray(signatureBytes)
        signatureBytes[0] += 1 + recovery_param
    print("here")
    #print(v0FieldsWithoutSignature)
    print(signatureLength)
    print(signatureBytes)
    print(v1FieldsBuffer)
    
    signed_transaction_bytes = (
    v0FieldsWithoutSignature
    + signatureLength
    + signatureBytes
    + v1FieldsBuffer
)

    return signed_transaction_bytes.hex()