import random

def H(x):
    """A stupid random-ish hash function.  In real life, this would be Blake2
    or AES.  Takes 2 32-bit inputs and returns a 32-bit output"""
    return (x*12345 + 7890) % (1 << 32)

def F(x):
    """F(x) should be a cryptographic pseudo-random permutation, which we can
    build from any cryptographic hash function.  In this case it is its own
    inverse, simplifying the code a bit."""
    l = x & 0xffffffff
    h = x >> 32
    h ^= H(l)
    l ^= H(h)
    h ^= H(l)
    l ^= H(h)
    h ^= H(l)
    return (h << 32) | l

def add(x, y):
    """This is what I am calling the TwoCats construction.  Generally it has the form:
           Finv(F(x) + F(y))"""
    tmp = F(x) + F(y)
    if tmp > 1 << 64:
        tmp -= 1 << 64
    return F(tmp)

def mul(m, x):
    r = x
    while m != 0:
        if m & 1:
            r = add(r, x)
        m >>= 1
        x = add(x, x)
    return r

sa = random.randrange(0, 1 << 64)
sb = random.randrange(0, 1 << 64)
print "Check that F is it's own inverse:", sa, "=", F(F(sa))
print "associative check on group addition law:"
print "(2 @ 3) @ 4 =", add(add(2, 3), 4)
print "2 @ (3 @ 4) =", add(2, add(3, 4))

g = 0
print "sa:", sa
print "sb:", sb
A = mul(sa, g)
B = mul(sb, g)
print "Alice public key:", A
print "Bob public key:", B
aliceMS = mul(sa, B)
bobMS   = mul(sb, A)
print "Alice master key:", aliceMS
print "Bob master key:", bobMS
assert aliceMS == bobMS