#!/usr/bin/env python2.7
"""
Self-contained Python ECDSA implementation for elliptic curves over prime
fields.

Jacob Welsh (PGP: 2599 9E50 95BA 9A2E 1CE6  305E A890 2A7C ACFD 13B0)

Version 20K, September 2016

Note that execution time for public key generation and signing is dependent on
the private key. NOT FOR USE IN SYSTEMS WHERE TIMING MAY BE OBSERVED BY
ADVERSARIES.
"""

# Let a point on an elliptic curve over the prime field F(p) be represented as
# either the pair (x, y) where x and y are the integer coordinates in [0, p),
# or Inf, the point at infinity.

class Inf(object):
    "Special value for the point at infinity (EC additive identity)"
    pass

def on_curve(point, p, a, b):
    """
    Test whether point is on the elliptic curve over the prime field of order p
    defined by:

        y^2 === x^3 + ax + b (mod p)

    ("===" is used herein to mean modular congruence.)
    """
    if point is Inf:
        return True
    x, y = point
    return (y**2 - (x**3 + a*x + b)) % p == 0

def ec_add(p, a):
    "Returns the EC group operation using parameters p, a as in on_curve"
    def add(p1, p2):
        "Add the points p1, p2"
        # Per [SEC1] section 2.2.1
        if p1 is Inf:
            return p2
        if p2 is Inf:
            return p1
        x1, y1 = p1
        x2, y2 = p2
        if x1 == x2:
            if y1 != y2 or y1 == 0:
                # Adding the negative point
                if y1 != 0:
                    assert y1 == p - y2, 'points off curve or out of bounds'
                return Inf
            # Adding the same point with nonzero y (doubling)
            slope = ((3*x1**2 + a) * mod_inverse(2*y1, p)) % p
        else:
            slope = ((y2 - y1) * mod_inverse(x2 - x1, p)) % p
        x3 = (slope**2 - x1 - x2) % p
        y3 = (slope*(x1 - x3) - y1) % p
        return (x3, y3)
    return add

def mod_inverse(a, m):
    """
    Multiplicative inverse of a modulo m:
    a * mod_inverse(a, m) === 1 (mod m)
    """
    assert m > 1
    if a < 0 or a >= m:
        a = a % m
    assert a != 0, 'not invertible (a=0)'
    b = m
    # Extended Euclidean algorithm: find x, y where ax + by = gcd(a, b).
    # Per [HAC] Algorithm 2.142 / 2.107
    x2 = 1
    x1 = 0
    y2 = 0
    y1 = 1
    while b:
        q, r = divmod(a, b)
        x = x2 - q*x1
        y = y2 - q*y1
        a = b
        b = r
        x2 = x1
        x1 = x
        y2 = y1
        y1 = y
    # (gcd, x, y) = (a, x2, y2)
    # If gcd is 1, it follows that ax === 1 (mod b)
    assert a == 1, 'not invertible (a and m not coprime)'
    if x2 < 0:
        x2 += m
    return x2

def scalar_mult(k, g, add):
    """
    Scalar multiplication: computes the equivalent of k repeated additions of
    group element g, using the group operation add(x, y), in O(log k) time.

    Uses the simple recursive method, so sys.getrecursionlimit() gives an upper
    bound on the number of bits in k, typically ~1000.
    """
    assert k > 0 # We could return the identity element for k=0, but we'd need
                 # to be told what it is.
    def mult(k, g):
        if k == 1:
            return g
        half = mult(k // 2, g)
        product = add(half, half)
        if k & 1:
            product = add(product, g)
        return product
    return mult(k, g)

class CurveParams(object):
    """
    Elliptic curve domain parameters over a prime field
    p, a, b: order of the field and curve coefficients, as in on_curve
    G: base point (generator)
    n: order of G
    h: cofactor (n*h = number of points on the curve)
    """
    def __init__(self, p, a, b, G, n, h):
        self.p, self.a, self.b = p, a, b
        self.G, self.n, self.h = G, n, h

def ec_priv_key(params, rand_bytes):
    """
    Generate private key (integer) for the curve given by params.
    rand_bytes(n) -> string of n random bytes (e.g. os.urandom)
    """
    return rand_int(params.n - 1, rand_bytes) + 1

def ec_pub_key(params, priv_key):
    "Compute public key (point) from priv_key for the curve given by params"
    return scalar_mult(priv_key, params.G, ec_add(params.p, params.a))

def ec_pub_key_valid(params, pub_key):
    "Check whether pub_key is valid for the curve given by params"
    # Per [SEC1] section 3.2.2.1
    p, a, b, n, h = params.p, params.a, params.b, params.n, params.h
    return (pub_key is not Inf and
            0 <= pub_key[0] < p and
            0 <= pub_key[1] < p and
            on_curve(pub_key, p, a, b) and
            (h == 1 or scalar_mult(n, pub_key, ec_add(p, a)) is Inf))

def ecdsa_sign(params, message_hash, priv_key, rand_bytes):
    """
    Return ECDSA signature (r, s) of message_hash using priv_key, for the curve
    given by params. Always succeeds, assuming valid inputs and properly
    functioning RNG.

    rand_bytes(n) -> string of n random bytes (e.g. os.urandom)

    As (r, -s mod n) is also a valid signature, the result is canonicalized to
    use the lesser of the two possible s-values.
    """
    # Per [SEC1] section 4.1.3
    n = params.n
    e = hash_to_int(message_hash, n.bit_length())
    while True:
        k = ec_priv_key(params, rand_bytes)
        R = ec_pub_key(params, k)
        r = R[0] % n
        if r == 0:
            continue
        s = (mod_inverse(k, n) * (e + r*priv_key)) % n
        if s == 0:
            continue
        return (r, min(s, n - s))

def ecdsa_verify(params, sig, message_hash, pub_key):
    """
    Checks whether the (r, s) pair sig is a valid ECDSA signature of
    message_hash using pub_key, for the curve given by params.
    """
    # Per [SEC1] section 4.1.4
    n = params.n
    r, s = sig
    if not ((0 < r < n) and (0 < s < n)):
        return False
    e = hash_to_int(message_hash, n.bit_length())
    s_inv = mod_inverse(s, n)
    u1 = (e*s_inv) % n
    u2 = (r*s_inv) % n
    add = ec_add(params.p, params.a)
    R = add(scalar_mult(u1, params.G, add), scalar_mult(u2, pub_key, add))
    if R is Inf:
        return False
    return R[0] % n == r

def bytes_to_int(b):
    "Convert big-endian byte sequence to unsigned integer"
    i = 0
    for byte in b:
        i = (i << 8) + ord(byte)
    return i

def hash_to_int(H, max_bits):
    """
    Convert the max_bits leftmost bits of big-endian byte sequence H to
    unsigned integer
    """
    # Equivalent to [SEC1] section 4.1.3 step 5
    e = bytes_to_int(H)
    hash_bits = 8*len(H)
    extra_bits = hash_bits - max_bits
    if extra_bits > 0:
        e = e >> extra_bits
    return e

def rand_int(n, rand_bytes):
    """
    Random integer in the interval [0, n) without modulo bias.
    rand_bytes(n) -> string of n random bytes (e.g. os.urandom)
    """
    if n <= 1:
        raise ValueError(n)
    nbytes, rem = divmod((n - 1).bit_length(), 8)
    if rem:
        nbytes += 1
    # Collect one more byte than strictly necessary to avoid cases where a
    # large part of the range is invalid (e.g. n=130)
    nbytes += 1
    rand_max = (1 << (8*nbytes)) - 1
    max_unbiased = rand_max - (rand_max % n)
    while True:
        r = bytes_to_int(rand_bytes(nbytes))
        if r <= max_unbiased:
            return r % n


##################################################
# Well-known curve parameters

# Generalized Koblitz curve over a 256-bit prime field, per [SEC2]
secp256k1 = CurveParams(
    p = 0xfffffffffffffffffffffffffffffffffffffffffffffffffffffffefffffc2f,
    a = 0,
    b = 7,
    G = (0x79be667ef9dcbbac55a06295ce870b07029bfcdb2dce28d959f2815b16f81798,
         0x483ada7726a3c4655da4fbfc0e1108a8fd17b448a68554199c47d08ffb10d4b8),
    n = 0xfffffffffffffffffffffffffffffffebaaedce6af48a03bbfd25e8cd0364141,
    h = 1)


##################################################
# Self-tests

def test_conversions():
    msg = '\x01\x02\x03\xfd\xfe\xff'
    assert bytes_to_int(msg) == 0x10203fdfeff
    assert hash_to_int(msg, 20) == 0x1020

def test_mod_inverse():
    modulus = 10007
    for a in xrange(1, modulus):
        a_inv = mod_inverse(a, modulus)
        assert 0 <= a_inv < modulus
        assert (a*a_inv) % modulus == 1

def test_scalar_mult():
    add = lambda x, y: x + y
    mul = lambda x, y: x*y
    modulus = 13
    mod_mul = lambda x, y: (x * y) % modulus
    for k in range(1, 20):
        for g in range(-20, 20):
            # Test on the group of integers under addition
            prod = scalar_mult(k, g, add)
            assert prod == k*g
            # Test on integers under multiplication (not actually a group,
            # since it's not invertible, but scalar_mult doesn't require that)
            prod = scalar_mult(k, g, mul)
            assert prod == g**k
        for g in range(0, modulus):
            # Test on the group of integers modulo a prime under multiplication
            prod = scalar_mult(k, g, mod_mul)
            assert prod == pow(g, k, modulus)

def validate_params(params, security_level):
    p, a, b = params.p, params.a, params.b
    G, n, h = params.G, params.n, params.h
    # EC parameter validation, per [SEC1] section 3.1.1.2.1
    try:
        from numbertheory import is_prime
        # e.g. http://webpages.charter.net/curryfans/peter/numbertheory.py
    except ImportError:
        print "warning: numbertheory not found; prime tests disabled"
        is_prime = lambda _: True
    assert p > 2
    assert is_prime(p)
    assert p.bit_length() >= 2*security_level
    for val in (a, b, G[0], G[1]):
        assert 0 <= val < p
    assert (4*a**3 + 27*b**2) % p != 0
    assert on_curve(G, p, a, b)
    assert is_prime(n)
    assert h <= 2**(security_level//8)
    assert h == int((p**0.5 + 1)**2 / n)
    assert scalar_mult(n, G, ec_add(p, a)) is Inf
    for b in range(1, 100):
        assert pow(p, b, n) != 1
    assert n != p

def test_keys(params):
    def mock_rand_bytes():
        def rand_byte_gen(x=0):
            m = 2**64
            while True:
                x = (6364136223846793005*x + 1) % m
                yield chr((x >> 32) & 255)
        gen = rand_byte_gen()
        return lambda n: ''.join(next(gen) for _ in range(n))
    rand_bytes = mock_rand_bytes()
    priv = ec_priv_key(params, rand_bytes)
    pub = ec_pub_key(params, priv)
    assert ec_pub_key_valid(params, pub)
    h = 'this is not a hash'
    sig = ecdsa_sign(params, h, priv, rand_bytes)
    assert ecdsa_verify(params, sig, h, pub)
    h = h*50
    r, s = ecdsa_sign(params, h, priv, rand_bytes)
    assert ecdsa_verify(params, (r, s), h, pub)
    assert not ecdsa_verify(params, (r+1, s), h, pub)
    assert not ecdsa_verify(params, (r, s+1), h, pub)

if __name__ == '__main__':
    print 'Testing byte conversions'
    test_conversions()
    print 'Testing group functions'
    test_mod_inverse()
    test_scalar_mult()
    print 'Verifying secp256k1 parameters'
    assert secp256k1.p == 2**256 - 2**32 - 2**9 - 2**8 - 2**7 - 2**6 - 2**4 - 1
    validate_params(secp256k1, 128)
    print 'Testing key generation, signing, and verification'
    test_keys(secp256k1)


##################################################
# Works cited
#
# [SEC1] Certicom Research 2009. "Standards for Efficient Cryptography: SEC 1:
#       Elliptic Curve Cryptography". Certicom Corp. Version 2.0.
#       http://www.secg.org/sec1-v2.pdf
#
# [SEC2] Certicom Research 2010. "Standards for Efficient Cryptography: SEC 2:
#       Recommended Elliptic Curve Domain Parameters". Certicom Corp. Version
#       2.0. http://www.secg.org/sec2-v2.pdf
#
# [HAC] Menezes, A., van Oorschot, P. and Vanstone, S. 1996. "Handbook of
#       Applied Cryptography". CRC Press. http://www.cacr.math.uwaterloo.ca/hac