from __future__ import annotations
from hashlib import new
from nacl.signing import SigningKey, SignedMessage, VerifyKey
from secrets import token_bytes
import nacl.bindings


"""
Demonstration of Twisted Edwards Curve 2^255-19 digital signature system, its
properties, and its adaptations. See below sources for more details.

Ed25519 (RFC 8032): https://www.rfc-editor.org/rfc/rfc8032.html
Adapters: https://medium.com/crypto-garage/adaptor-signature-schnorr-signature-and-ecdsa-da0663c2adc4
AMHLs: https://eprint.iacr.org/2018/472
"""

# 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 nacl.bindings.crypto_core_ed25519_scalar_reduce(H_big(*parts))

def derive_key_from_seed(seed: bytes) -> bytes:
    """Derive the scalar used for signing from a seed."""
    return clamp_scalar(H_big(seed)[:32], True)

def aggregate_points(points: list) -> bytes:
    """Aggregate points on the Ed25519 curve."""
    # type checking inputs
    for pt in points:
        if type(pt) is not bytes and type(pt) is not VerifyKey:
            raise TypeError('each point must be bytes or VerifyKey')

    # normalize points to bytes
    points = [pt if type(pt) is bytes else bytes(pt) for pt in points]

    # raise an error for invalid points
    for pt in points:
        if not nacl.bindings.crypto_core_ed25519_is_valid_point(pt):
            raise ValueError('each point must be a valid ed25519 point')

    # compute the sum
    sum = points[0]
    for i in range(1, len(points)):
        sum = nacl.bindings.crypto_core_ed25519_add(sum, points[i])

    return sum

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 hexify(thing):
    """Convert a data structure to something printable."""
    if type(thing) is dict:
        new_thing = {}
        for key in thing:
            value = thing[key]
            key = key.hex() if type(key) is bytes else key
            new_thing[key] = hexify(value)
        return new_thing

    if type(thing) is list:
        return [hexify(t) for t in thing]

    if type(thing) is tuple:
        return tuple([hexify(t) for t in thing])

    if type(thing) is bytes:
        return thing.hex()

    return thing


# adapter functions
def public_tweak_adapter(seed: bytes, m: bytes, T: bytes) -> tuple[bytes]:
    """Create a public tweak adapter signature. Return vbalue is (R, T, sa)."""
    x = derive_key_from_seed(seed)
    X = nacl.bindings.crypto_scalarmult_ed25519_base_noclamp(x) # G^x
    nonce = H_big(seed)[32:]
    r = clamp_scalar(H_small(H_big(nonce, m))) # H(nonce || m)
    R = nacl.bindings.crypto_scalarmult_ed25519_base_noclamp(r) # G^r
    RT = aggregate_points((R, T)) # R + t
    ca = clamp_scalar(H_small(RT, X, m)) # H(R + T || X || m)
    sa = nacl.bindings.crypto_core_ed25519_scalar_add(r, nacl.bindings.crypto_core_ed25519_scalar_mul(ca, x)) # r + H(R + T || X || m) * x
    return (R, T, sa)

def verify_public_tweak_adapter(adapter: tuple[bytes], m: bytes, X: bytes) -> bool:
    """Verify a public tweak adapter is valid for a given message."""
    (R, T, sa) = adapter
    sa_G = nacl.bindings.crypto_scalarmult_ed25519_base_noclamp(sa) # sa_G = G^sa
    RT = aggregate_points((R, T)) # R + T
    ca = clamp_scalar(H_small(RT, X, m)) # H(R + T || X || m)
    caX = nacl.bindings.crypto_scalarmult_ed25519_noclamp(ca, X) # X^H(R + T || X || m)
    RcaX = aggregate_points((R, caX)) # R + X^H(R + T || X || m)
    return bytes_are_same(sa_G, RcaX) # G^sa == R + X^H(R + T || X || m)

def decrypt_public_tweak_adapter(adapter: tuple[bytes], m: bytes, t: bytes) -> SignedMessage:
    """Decrypt a public tweak adapter signature."""
    (R, T, sa) = adapter
    RT = aggregate_points((R, T)) # R + T
    s = nacl.bindings.crypto_core_ed25519_scalar_add(sa, t) # s = sa + t
    sig = SignedMessage(RT + s + m)
    return sig


# AMHL functions
# Messy conventions taken directly from the Anonymous Multi-Hop Lock paper
def sample() -> bytes:
    """Take 1 sample from the domain of the homomorphic one-way function."""
    return clamp_scalar(token_bytes(32))

def samples(n: int) -> tuple[bytes]:
    """Take n samples from the domain of the homomorphic one-way function."""
    return tuple(sample() for i in range(n))

def oneway(scalar: bytes) -> bytes:
    """Run the homomorphic one-way function on the input scalar."""
    return nacl.bindings.crypto_scalarmult_ed25519_base_noclamp(scalar)

def setup(n_users: int) -> tuple[tuple[bytes], tuple[bytes]]:
    """Setup the lock inputs."""
    y = samples(n_users)
    Y = [oneway(y[0])]
    for i, y_i in enumerate(y):
        if i > 0:
            Y.append(aggregate_points((Y[i-1], oneway(y_i))))
    return (y, tuple(Y))

def scalar_sum(*scalars: tuple[bytes]) -> bytes:
    """Compute the sum of the scalars."""
    sum = scalars[0]
    for i in range(1, len(scalars)):
        sum = nacl.bindings.crypto_core_ed25519_scalar_add(sum, scalars[i])
    return sum

def setup_for(s: tuple[tuple[bytes], tuple[bytes]], i: int) -> tuple[bytes]:
    """Generate the setup for a particular user given the setup s and
        the user index i.
    """
    if i == 0:
        return (s[0][0],)

    if i == len(s[0]):
        return ((s[1][i-1], 0, 0), scalar_sum(*s[0]))

    return (s[1][i-1], s[1][i], s[0][i])

def check_setup(s: tuple[bytes], i: int, n: int) -> bool:
    """Verifies the setup for the ith of n users is valid."""
    if i == 0:
        return len(s) == 1 and isinstance(s[0], bytes)

    if i == n:
        return len(s) == 2 and type(s[0]) is tuple and len(s[0]) == 3 and type(s[0][0]) is bytes and type(s[1]) is bytes

    Y_i = aggregate_points((s[0], oneway(s[2])))
    return bytes_are_same(Y_i, s[1])

def lock(s: tuple) -> tuple[bytes, bool]:
    """Create the lock from the setup."""
    return (s[1], False)

def release(k: bytes, sI: tuple) -> bytes:
    """Release a left lock given a key from a released intermediate lock."""
    y = sI[2] if type(sI) is tuple else sI
    return nacl.bindings.crypto_core_ed25519_scalar_sub(k, y)

def verify_lock_key(l: bytes, k: bytes) -> bool:
    """Verify that a key opens a lock."""
    return bytes_are_same(l, oneway(k))


# tests
def test_sign_and_verify():
    """Replication of the maths involved in signing and verifying."""
    # set seed and generate key pair from it
    seed = token_bytes(32)
    skey = SigningKey(seed)

    # derive key pair manually
    x = derive_key_from_seed(seed) # clamp(H(seed)[:32])
    X = nacl.bindings.crypto_scalarmult_ed25519_base_noclamp(x) # G^x
    m = b'hello world'

    # create signature with nacl
    sig = skey.sign(m)

    # create signature manually
    nonce = H_big(seed)[32:]
    r = clamp_scalar(H_small(nonce, m))
    R = nacl.bindings.crypto_scalarmult_ed25519_base_noclamp(r) # G^r
    c = clamp_scalar(H_small(R, X, m)) # clamp(H(R || X || m))
    s = nacl.bindings.crypto_core_ed25519_scalar_add(r, nacl.bindings.crypto_core_ed25519_scalar_mul(c, x)) # r + c*x
    nsig = SignedMessage(R + s + m)

    # G, X, R, s, and m are public values
    # c can be derived: H(R || X || m)
    # Therefore, G^s = R * X^c can be verified by anyone
    # But only the holder of x and r can create the signature

    # verify both
    assert bytes_are_same(bytes(sig), bytes(nsig))
    assert skey.verify_key.verify(sig) # s = c*x + r
    assert skey.verify_key.verify(nsig) # G^s = G^(c*x + r) = R * G^(x*c) = R * X^c


def test_homomorphic_one_way():
    """Test if Ed25519 meets the homomorphic one way condition."""
    x1, x2 = clamp_scalar(token_bytes(32)), clamp_scalar(token_bytes(32))
    y1 = nacl.bindings.crypto_scalarmult_ed25519_base_noclamp(x1) # G^x1
    y2 = nacl.bindings.crypto_scalarmult_ed25519_base_noclamp(x2) # G^x2

    # test
    y3_1 = nacl.bindings.crypto_core_ed25519_add(y1, y2) # G^x1 + G^x2
    x3 = nacl.bindings.crypto_core_ed25519_scalar_add(x1, x2) # x1 + x2
    y3_2 = nacl.bindings.crypto_scalarmult_ed25519_base_noclamp(x3) # G^(x1+x2)
    assert y3_1 == y3_2 # G^x1 * G^x2 = G^(x1+x2) where * denotes group operator


def test_adapter_signature():
    """Test the construction, unlocking, and verification of adapters.
        For an explanation and possible uses, see this medium article:
        https://medium.com/crypto-garage/adaptor-signature-schnorr-signature-and-ecdsa-da0663c2adc4
    """
    # key pairs and message
    seed1, seed2 = token_bytes(32), token_bytes(32)
    skey = SigningKey(seed1)
    x, t = derive_key_from_seed(seed1), derive_key_from_seed(seed2)
    X = nacl.bindings.crypto_scalarmult_ed25519_base_noclamp(x) # G^x
    T = nacl.bindings.crypto_scalarmult_ed25519_base_noclamp(t) # G^t
    m = b'txn sending money to counterparty as part of payment route'

    # private key tweaking construction
    # sa = t + r + H(R || X || m) * x
    nonce = H_big(seed1)[32:]
    r = clamp_scalar(H_small(nonce, m))
    R = nacl.bindings.crypto_scalarmult_ed25519_base_noclamp(r) # G^r
    c = clamp_scalar(H_small(R, X, m)) # clamp(H(R || X || m))
    tr = nacl.bindings.crypto_core_ed25519_scalar_add(t, r)
    sa = nacl.bindings.crypto_core_ed25519_scalar_add(tr, nacl.bindings.crypto_core_ed25519_scalar_mul(c, x)) # t + r + c*x

    # public values: R, sa, m, X, T
    # sig = (T, R, sa, m)

    # private key tweaking adapter verification
    # sa_G = G^sa
    sa_G = nacl.bindings.crypto_scalarmult_ed25519_base_noclamp(sa)
    c = clamp_scalar(H_small(R, X, m)) # clamp(H(R || X || m))
    cX = nacl.bindings.crypto_scalarmult_ed25519_noclamp(c, X)
    TRcX = aggregate_points((T, R, cX))
    assert bytes_are_same(sa_G, TRcX)
    # G^sa == T + R + X^H(R || X || m)

    # private key tweaking signature decryption and verification
    # s = sa - t
    s = nacl.bindings.crypto_core_ed25519_scalar_sub(sa, t)
    sig = SignedMessage(R + s + m)
    assert skey.verify_key.verify(sig)
    # G^(sa - t) == T + R + X^H(R || X || m) - T
    # G^s == R + X^H(R || X || m)


    # public key tweaking construction
    # sa = r + H(R + T || X || m) * x
    RT = aggregate_points((R, T))
    ca = clamp_scalar(H_small(RT, X, m))
    sa = nacl.bindings.crypto_core_ed25519_scalar_add(r, nacl.bindings.crypto_core_ed25519_scalar_mul(ca, x))

    # public values: R, sa, m, X, T
    # ca and RT can be derived from public values
    # sig = (T, R, sa, m)

    # public key tweaking adapter verification
    # sa_G = G^sa
    sa_G = nacl.bindings.crypto_scalarmult_ed25519_base_noclamp(sa)
    caX = nacl.bindings.crypto_scalarmult_ed25519_noclamp(ca, X)
    RcaX = aggregate_points((R, caX))
    assert bytes_are_same(sa_G, RcaX)
    # G^sa == R + X^H(R + T || X || m)

    # public key tweaking signature decryption and verification
    # s = sa + t
    s = nacl.bindings.crypto_core_ed25519_scalar_add(sa, t)
    sig = SignedMessage(RT + s + m)
    assert skey.verify_key.verify(sig)
    # G^(sa + t) == R + X^H(R + T || X || m) + T
    # G^s == (R + T) + X^H((R + T) || X || m)
    ...


def test_AMHL():
    """Test for setup, locking, and release of an Anonymous Multi-Hop
        Lock using the homomorphic qualities of ed25519.
    """
    # first run the initial setup for 5 payers (4 intermediate)
    n = 5
    s = setup(n)

    # validate setups for each user
    for i in range(len(s[0])):
        assert check_setup(setup_for(s, i), i, n)

    # validate releasing of locks from right to left
    s_n = setup_for(s, n)
    k = s_n[1]

    for i in range(n-1, 1, -1):
        s_i = setup_for(s, i)
        r = release(k, s_i)
        assert verify_lock_key(setup_for(s, i-1)[1], r)
        k = r


def test_AMHL_adapters():
    """Test for combining the AMHL primitive with the adapter signature
        primitive for atomic transaction locking/unlocking. In this case,
        Alice pays Dave through Bob and Carla, with Bob and Carla each
        taking a fee for faciliating the transaction.
    """
    # first run the initial setup for 3 payers (2 intermediate)
    n = 3
    s = setup(n)

    # validate setups for each user
    for i in range(len(s[0])):
        assert check_setup(setup_for(s, i), i, n)

    # Alice setup
    Alice = {
        'seed' : token_bytes(32),
        'outbound_txn' : {
            'm': b'Alice pays Bob 12'
        },
        'lock': setup_for(s, 0)
    }
    Alice['skey'] = SigningKey(Alice['seed'])
    Alice['vkey'] = Alice['skey'].verify_key
    Alice['x'] = derive_key_from_seed(Alice['seed'])
    Alice['X'] = nacl.bindings.crypto_scalarmult_ed25519_base_noclamp(Alice['x'])
    # create adapter
    Alice['outbound_txn']['adapter'] = public_tweak_adapter(Alice['seed'], Alice['outbound_txn']['m'], oneway(Alice['lock'][0]))

    # Bob setup
    Bob = {
        'seed' : token_bytes(32),
        'inbound_txn': Alice['outbound_txn'],
        'outbound_txn' : {
            'm': b'Bob pays Carla 11'
        },
        'lock': setup_for(s, 1)
    }
    Bob['skey'] = SigningKey(Bob['seed'])
    Bob['vkey'] = Bob['skey'].verify_key
    Bob['x'] = derive_key_from_seed(Bob['seed'])
    Bob['X'] = nacl.bindings.crypto_scalarmult_ed25519_base_noclamp(Bob['x'])
    # verify Bob's left lock is Alice's right lock
    assert bytes_are_same(Bob['lock'][0], oneway(Alice['lock'][0]))
    # verify Alice's adapter
    assert verify_public_tweak_adapter(Bob['inbound_txn']['adapter'], Bob['inbound_txn']['m'], Alice['X'])
    # create adapter
    Bob['outbound_txn']['adapter'] = public_tweak_adapter(Bob['seed'], Bob['outbound_txn']['m'], Bob['lock'][1])

    # Carla setup
    Carla = {
        'seed' : token_bytes(32),
        'inbound_txn' : Bob['outbound_txn'],
        'outbound_txn': {
            'm': b'Carla pays Dave 10'
        },
        'lock': setup_for(s, 2)
    }
    Carla['skey'] = SigningKey(Carla['seed'])
    Carla['vkey'] = Carla['skey'].verify_key
    Carla['x'] = derive_key_from_seed(Carla['seed'])
    Carla['X'] = nacl.bindings.crypto_scalarmult_ed25519_base_noclamp(Carla['x'])
    # verify Carla's left lock is Bob's right lock
    assert bytes_are_same(Carla['lock'][0], Bob['lock'][1])
    # verify Bob's adapter
    assert verify_public_tweak_adapter(Carla['inbound_txn']['adapter'], Carla['inbound_txn']['m'], Bob['X'])
    # create adapter
    Carla['outbound_txn']['adapter'] = public_tweak_adapter(Carla['seed'], Carla['outbound_txn']['m'], Carla['lock'][1])

    # Dave setup
    Dave = {
        'seed' : token_bytes(32),
        'inbound_txn' : Carla['outbound_txn'],
        'lock': setup_for(s, 3)
    }
    # verify Dave's left lock is Carla's right lock
    assert bytes_are_same(Dave['lock'][0][0], Carla['lock'][1])
    # verify Carla's adapter
    assert verify_public_tweak_adapter(Dave['inbound_txn']['adapter'], Dave['inbound_txn']['m'], Carla['X'])

    # decrypt adapter paying Dave
    k = Dave['lock'][1]
    Dave['inbound_txn']['signature'] = decrypt_public_tweak_adapter(Dave['inbound_txn']['adapter'], Dave['inbound_txn']['m'], k)
    assert Carla['vkey'].verify(Dave['inbound_txn']['signature'])

    # release lock for Carla
    r = release(k, Carla['lock'])
    assert verify_lock_key(Carla['lock'][0], r) # Carla's left lock == Bob's right lock

    # decrypt adapter paying Carla
    Carla['inbound_txn']['signature'] = decrypt_public_tweak_adapter(Carla['inbound_txn']['adapter'], Carla['inbound_txn']['m'], r)
    assert Bob['vkey'].verify(Carla['inbound_txn']['signature'])

    # release lock for Bob
    r = release(r, Bob['lock'])
    assert verify_lock_key(Bob['lock'][0], r) # Bob's left lock == Alice's right lock

    # decrypt adapter paying Bob
    Bob['inbound_txn']['signature'] = decrypt_public_tweak_adapter(Bob['inbound_txn']['adapter'], Bob['inbound_txn']['m'], r)
    assert Alice['vkey'].verify(Bob['inbound_txn']['signature'])

    # Dave now has a signed transaction paying him from Carla
    # Carla now has a signed transaction paying her from Bob
    # Bob now has a signed transaction paying him from Alice
    # Alice has paid Dave through Bob and Carla
    ...


def license():
    """Copyleft (c) 2022 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__


if __name__ == '__main__':
    for i in range(128):
        test_homomorphic_one_way()
    print('homomorphic one-way function test passed')

    for i in range(128):
        test_sign_and_verify()
    print('signature creation and verification test passed')

    for i in range(128):
        test_adapter_signature()
    print('adapter signature test passed')

    for i in range(128):
        test_AMHL()
    print('anonymous multihop lock test passed')

    for i in range(128):
        test_AMHL_adapters()
    print('AMHL adapter (payment channel routing) test passed')