"""This module implements SRP-6a (see http://srp.stanford.edu/design.html).
Variable names used here match those used in the description on that page:

    N       a large prime number (N = 2q + 1, where q is prime)
    g       a generator (primitive root) modulo N
    k       a multiplier parameter (k = H(N, g) in SRP-6a)
    s       a random string used as the user's salt
    I       the username
    p       the user's cleartext password
    H()     a one-way hash function
    u       a random scrambling parameter (public)
    a, b    ephemeral private keys, generated randomly and kept secret
    A, B    corresponding ephemeral public keys
    x       a private key derived from p and s
    v       the host's password verifier
    K       the session key

The host stores passwords using the following formula:

    x = H(s, p)               (s is chosen randomly)
    v = g^x                   (computes password verifier)

The host then keeps {I, s, v} in its password database. The authentication
protocol itself goes as follows:

User -> Host: I, A = g^a                  (identifies self, a = random number)
Host -> User: s, B = kv + g^b             (sends salt, b = random number)

        Both: u = H(A, B)

        User: x = H(s, p)                 (user enters password)
        User: S = (B - kg^x) ^ (a + ux)   (computes session key)
        User: K = H(S)

        Host: S = (Av^u) ^ b              (computes session key)
        Host: K = H(S)

Now the two parties have a shared, strong session key K. To complete
authentication, they need to prove to each other that their keys match.
One possible way:

User -> Host:  M = H(H(N) xor H(g), H(I), s, A, B, K)
Host -> User:  H(A, M, K)

The two parties also employ the following safeguards:

    1.  The user will abort if he receives B == 0 (mod N) or u == 0.
    2.  The host will abort if it detects that A == 0 (mod N).
    3.  The user must show his proof of K first. If the server detects
        that the user's proof is incorrect, it must abort without showing
        its own proof of K. 
"""

__author__ = 'Ka-Ping Yee'
__date__ = '2006-06-25'

from random import randrange

class SRPError(Exception): pass

class Object:
    def __init__(self, *args, **kw):
        self.__dict__ = kw
        for arg in args:
            self.__dict__[arg.__name__] = arg

def SRP(N, g, hash, _pow=pow):
    def value(str):
        value = 0L
        for ch in str:
            value = (value << 8L) + ord(ch)
        return value

    def H(*values):
        return value(hash('-'.join(map(str, values)))) % N

    def pow(base, exponent):
        return _pow(base, exponent, N)

    def exp(exponent):
        return pow(g, exponent)

    k = H(N, g)

    def random():
        return randrange(2, N)

    def ensure(condition, exception):
        if not condition:
            raise SRPError(exception)

    def setup(password):
        p = value(password)
        s = random()
        x = H(s, p)
        v = exp(x)
        return s, v

    # Each side of the protocol is implemented as a generator function.
    # The first thing to be yielded is always the push() function.
    # yield <value> means "send a value to the other side"
    # yield None means "wait for a value from the other side to be push()ed"

    def login(username, password, R):
        I = value(username) % N
        p = value(password) % N

        # First emit the function by which we will receive messages.
        queue = []
        yield queue.append

        yield 1
        a = random()
        A = exp(a)
        yield A
        yield None; s = queue.pop(0) % N
        yield None; B = queue.pop(0) % N
        ensure(B, 'received zero for B')
        u = H(A, B)
        ensure(u, 'received zero for u')
        x = H(s, p)
        S = pow(B - k*exp(x), a + u*x)
        K = H(S)
        M = H(H(N)^H(g), H(I), s, A, B, K)
        yield M
        yield None; remote_HAMK = queue.pop(0)
        ensure(remote_HAMK == H(A, M, K), 'H(A, M, K) did not match')
        R.key = K

    def authenticate(username, s, v, R):
        I = value(username) % N
        s = s % N
        v = v % N

        # First emit the function by which we will receive messages.
        queue = []
        yield queue.append

        yield None; protocol = queue.pop(0)
        if protocol == 1:
            yield None; A = queue.pop(0) % N
            ensure(A, 'received zero for A')
            b = random()
            B = (k*v + exp(b)) % N
            yield s
            yield B
            u = H(A, B)
            S = pow(A*pow(v, u), b)
            K = H(S)
            M = H(H(N)^H(g), H(I), s, A, B, K)
            yield None; remote_M = queue.pop(0)
            ensure(remote_M == M, 'M did not match')
            yield H(A, M, K)
            R.key = K
        else:
            raise SRPError('protocol %d is unsupported' % protocol)

    return Object(setup, login, authenticate)