"""Implementation of GCM (Galois/Counter Mode).  See http://csrc.nist.gov/
CryptoToolkit/modes/proposedmodes/gcm/gcm-revised-spec.pdf ."""

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

def blockxor(a, b):
    return [x ^ y for x, y in zip(a, b)]

def mult(X, Y):
    """Multiply in GF(2^128) modulo x^128 + x^7 + x^2 + x + 1."""
    Z, V = 0, X
    bit = 1<<127
    R = (bit>>7) + (bit>>2) + (bit>>1) + (bit>>0)
    for i in range(128):
        if Y & (bit>>i):
            Z ^= V
        if V & 1 == 0:
            V = (V>>1)
        else:
            V = (V>>1) ^ R
    return Z

def byte(n):
    return n % (1<<8)

def tolong(block):
    """Convert a block (list of integer bytes) to a long integer."""
    return sum([b<<((len(block) - 1 - i)*8) for i, b in enumerate(block)])

def toblock(long, blockbytes):
    """Convert a long integer to a block (list of integer bytes)."""
    return [byte(long>>shift) for shift in range(blockbytes*8 - 8, -8, -8)]

def toblocks(data, blockbytes):
    """Convert a string to a list of blocks (a block is a list of bytes)."""
    while len(data) % blockbytes != 0:
        data += '\x00'
    return [map(ord, data[i:i + blockbytes])
            for i in range(0, len(data), blockbytes)]

def tostring(blocks):
    """Convert a list of blocks to a string."""
    return ''.join([''.join(map(chr, block)) for block in blocks])

def GCM(cipher, iv=None, nonce=None):
    """Create a Galois/Counter Mode object with the given cipher.
    Supply a 96-bit long integer for the 'iv' argument or a single block
    as a long integer for the 'nonce' argument (which will be used to
    generate the IV).  The resulting object has a method encrypt() that
    takes a string and yields an encrypted string concatenated with a
    16-byte MAC, and a method decrypt() that takes an encrypted string
    with a 16-byte MAC suffix and yields the decrypted string."""

    if cipher.blockbytes != 16:
        raise TypeError('GCM requires a cipher with 128-bit blocks')

    if iv is None:
        if nonce is None:
            raise TypeError('either iv or nonce must be given')
        iv = tolong(cipher.encipher(toblock(nonce, cipher.blockbytes)))
    iv = iv % (1<<96)
    H = tolong(cipher.encipher(toblock(0, 16)))
    Y0 = toblock((iv<<32) | 1, 16)

    def transform(data):
        Y = Y0
        newblocks = []
        for block in toblocks(data, 16):
            Y = Y[:-4] + toblock(tolong(Y[-4:]) + 1, 4)
            newblocks.append(blockxor(block, cipher.encipher(Y)))
        return tostring(newblocks)[:len(data)]

    def ghash(data):
        X = 0
        for block in toblocks(data, 16):
            X = mult(X ^ tolong(block), H)
        return mult(X ^ (len(data)*8), H)

    def encrypt(data):
        C = transform(data)
        T = ghash(C) ^ tolong(cipher.encipher(Y0))
        return C, tostring([toblock(T, 16)])

    def decrypt(data, mac):
        T = tolong(map(ord, mac))
        if ghash(data) ^ tolong(cipher.encipher(Y0)) != T:
            raise ValueError('invalid message authentication code')
        return transform(data)

    return Object(encrypt, decrypt)
