vrf.py

"""Ed25519 verifiable random function. Important functions are prove,
proof2hash, and verify.
Reference used: https://datatracker.ietf.org/doc/html/draft-irtf-cfrg-vrf-02
NB: liberties were taken to avoid flaws in the hash_to_curve specification
documented in Section 5.4.1.1 and to slightly simplify the process over all;
instead of the computation time for hash_to_curve being variable and determined
by alpha, opening the process to timing attacks, this implementation does it in
one try by multiplying the public key by the hash of the input alpha.
"""


from hashlib import new
from nacl.signing import SigningKey, VerifyKey
from nacl.bindings import (
    crypto_core_ed25519_scalar_reduce,
    crypto_scalarmult_ed25519_noclamp,
    crypto_scalarmult_ed25519,
    crypto_scalarmult_ed25519_base_noclamp,
    crypto_core_ed25519_scalar_mul,
    crypto_core_ed25519_scalar_add,
    crypto_core_ed25519_sub,
)
from sys import argv


"""Helper functions"""
def clamp_scalar(scalar: bytes, from_private_key: bool = False) -> bytes:
    """Make a clamped scalar."""
    if type(scalar) is bytes and len(scalar) >= 32:
        x_i = bytearray(scalar[:32])
    elif type(scalar) is SigningKey:
        x_i = bytearray(new('sha512', bytes(scalar)).digest()[:32])
        from_private_key = True
    else:
        raise ValueError('not a SigningKey and not 32+ bytes scalar')

    if from_private_key:
        # set bits 0, 1, and 2 to 0
        # nb: lsb is right-indexed
        x_i[0] &= 0b11111000
        # set bit 254 to 1
        x_i[31] |= 0b01000000

    # set bit 255 to 0
    x_i[31] &= 0b01111111

    return bytes(x_i)

def H_big(*parts) -> bytes:
    """The big, 64-byte hash function."""
    return new('sha512', b''.join(parts)).digest()

def H_small(*parts) -> bytes:
    """The small, 32-byte hash function."""
    return crypto_core_ed25519_scalar_reduce(H_big(*parts))

def xor(b1: bytes, b2: bytes) -> bytes:
    """XOR two equal-length byte strings together."""
    b3 = bytearray()
    for i in range(len(b1)):
        b3.append(b1[i] ^ b2[i])

    return bytes(b3)

def bytes_are_same(b1: bytes, b2: bytes) -> bool:
    """Timing-attack safe bytes comparison."""
    return len(b1) == len(b2) and int.from_bytes(xor(b1, b2), 'little') == 0

def hash_to_curve(y: VerifyKey, alpha: bytes) -> bytes:
    """Return a curve point from y multiplied by the hash of alpha."""
    h = clamp_scalar(H_small(alpha))
    return crypto_scalarmult_ed25519_noclamp(h, bytes(y))


"""Main functions."""
def prove(sk: SigningKey, alpha: bytes):
    """Generate a proof."""
    if type(sk) is not SigningKey:
        raise TypeError('sk must be a nacl.signing.SigningKey')
    if type(alpha) is not bytes:
        raise TypeError('alpha must be bytes')

    xseed = H_big(bytes(sk))
    x = clamp_scalar(xseed[:32], True)

    h = hash_to_curve(sk.verify_key, alpha)
    gamma = crypto_scalarmult_ed25519(x, h)
    k = clamp_scalar(new('sha256', xseed[32:] + h).digest())
    gk = crypto_scalarmult_ed25519_base_noclamp(k)
    hk = crypto_scalarmult_ed25519_noclamp(k, h)
    c = H_small(h, gamma, gk, hk)
    cx = crypto_core_ed25519_scalar_mul(c, x)
    s = crypto_core_ed25519_scalar_add(k, cx)

    return (gamma, c, s)

def proof2hash(pi: tuple):
    """Get the hash of the proof."""
    if type(pi) is not tuple:
        raise TypeError('pi must be tuple of 3 values of 32 bytes each')
    if type(pi[0]) is not bytes:
        raise TypeError('pi must be tuple of 3 values of 32 bytes each')
    if len(pi[0]) != 32:
        raise ValueError('pi must be tuple of 3 values of 32 bytes each')

    three = b'\x03'
    (gamma, c, s) = pi
    cofactor = (8).to_bytes(32, 'little')
    gammacofactor = crypto_scalarmult_ed25519_noclamp(cofactor, gamma)
    return H_small(three, gammacofactor)

def verify(y: VerifyKey, alpha: bytes, pi: tuple):
    """Verify that the proof is valid."""
    if type(y) is not VerifyKey:
        raise TypeError('y must be a nacl.signing.VerifyKey')
    if type(alpha) is not bytes:
        raise TypeError('alpha must be bytes')
    if type(pi) is not tuple:
        raise TypeError('pi must be tuple of 3 values of 32 bytes each')
    if len(pi) != 3:
        raise ValueError('pi must be tuple of 3 values of 32 bytes each')
    for part in pi:
        if type(part) is not bytes or len(part) != 32:
            raise ValueError('pi must be tuple of 3 values of 32 bytes each')

    (gamma, c, s) = pi
    gs = crypto_scalarmult_ed25519_base_noclamp(s)
    yc = crypto_scalarmult_ed25519_noclamp(c, bytes(y))
    u = crypto_core_ed25519_sub(gs, yc)
    h = hash_to_curve(y, alpha)
    hs = crypto_scalarmult_ed25519_noclamp(s, h)
    gammac = crypto_scalarmult_ed25519_noclamp(c, gamma)
    v = crypto_core_ed25519_sub(hs, gammac)
    cprime = H_small(h, gamma, u, v)
    return bytes_are_same(c, cprime)


"""License."""
def license() -> str:
    """ISC License

        Copyleft (c) 2021, k98kurz

        Permission to use, copy, modify, and/or distribute this software for any
        purpose with or without fee is hereby granted, provided that the above
        copyleft notice and this permission notice appear in all copies.

        THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
        WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
        MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
        ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
        WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
        ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
        OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
    """
    return license.__doc__


"""Instantiation from the cli."""
def main(name, mode = '', value1 = '', value2 = '', value3 = ''):
    if mode == 'prove':
        if len(value1) != 64 or len(value2) < 1:
            return usage(name)
        sk = SigningKey(bytes.fromhex(value1))
        if value2[:4] == 'hex:':
            alpha = bytes.fromhex(value2[4:])
        else:
            alpha = bytes(value2, 'utf-8')
        pi = prove(sk, alpha)
        print(''.join([part.hex() for part in pi]))
    elif mode == 'verify':
        if len(value1) != 64 or len(value2) < 1 or len(value3) != 192:
            return usage(name)
        vk = VerifyKey(bytes.fromhex(value1))
        if value2[:4] == 'hex:':
            alpha = bytes.fromhex(value2[4:])
        else:
            alpha = bytes(value2, 'utf-8')
        proof = bytes.fromhex(value3)
        pi = (proof[:32], proof[32:64], proof[64:])
        print('valid' if verify(vk, alpha, pi) else 'invalid')
    elif mode == 'gethash':
        if len(value1) != 192:
            return usage(name)
        proof = bytes.fromhex(value1)
        pi = (proof[:32], proof[32:64], proof[64:])
        print(proof2hash(pi).hex())
    else:
        usage(name)

def usage(name):
    print(f'usage:\t{name} prove [skey_hex] [utf8string|hex:hexstring]')
    print(f'\t{name} verify [vkey_hex] [utf8string|hex:hexstring] [proof]')
    print(f'\t{name} gethash [proof]')


if __name__ == '__main__':
    main(*argv)