# Hack implementation of DJB's Poly1305-AES MAC.
# Written 2005-01-18 by Ken Raeburn, and placed in the public domain.
# Apologies for the clunkiness, I'm still learning Python.

from aes import aesprim
# Rather than picking apart the strings in C and doing bounds checking etc,
# I'm passing sequences and letting PyArgParseTuple figure out if they're
# the right type etc.  In retrospect, I should've just looked up how
# to raise an exception. :-)
def aes_encrypt(key, val):
  return aesprim(map(ord,key), map(ord,val))

aes_module_source = """
/* gcc -I${path}/include -O -c aes.c
   gcc -shared -o aes.so aes.o -L${path}/lib -lcrypto */
#include <python2.3/Python.h>
#include <openssl/aes.h>
static PyObject *aesprim (PyObject *self, PyObject *args) {
    unsigned char key[16], val[16], out[16]; AES_KEY k = { 0 };
#define BREAKOUT(X) &X[0], &X[1], &X[2], &X[3], &X[4], &X[5], &X[6], &X[7], \
		   &X[8], &X[9], &X[10], &X[11], &X[12], &X[13], &X[14], &X[15]
    if (!PyArg_ParseTuple(args, "(BBBBBBBBBBBBBBBB)(BBBBBBBBBBBBBBBB)",
			  BREAKOUT(key), BREAKOUT(val))) return NULL;
    assert (AES_set_encrypt_key(key, 16 * 8, &k) == 0);
    AES_encrypt(val, out, &k);
    return Py_BuildValue("s#", out, 16);
}
static PyMethodDef methods[] = {
    { "aesprim", aesprim, METH_VARARGS, "AES encryption primitive." }, { 0 } };
PyMODINIT_FUNC initaes(void) { (void) Py_InitModule("aes", methods); }
"""

# Hex strings are used in specifying the test cases to ensure that the
# byte order matches the examples, but the algorithm says they're treated
# as little-endian numbers.
# The [::-1] bit is new in Python 2.3; it reverses the string.
# Hexlify converts each byte to two hex digits; int(..., 16) reads a big
# number as hex.
import binascii
def str2num_littleend(val):
  return int(binascii.hexlify(val[::-1]), 16)
# Convert byte string to hex, separated by spaces, as in djb's paper.
def hexify(s):
  return ' '.join(map(binascii.hexlify, s))

def poly1305aes(k, r, n, msg):
  mod1305 = (1L << 130) - 5
  print "test case:"
  print "m        = ", hexify(msg)
  rval = str2num_littleend(r)
  q = (len(msg) + 15) / 16
  tot = 0
  for i in range(q):
    sub = msg[i*16 : i*16+16] + "\x01"
    sub += (17 - len(sub)) * "\x00"
    num = str2num_littleend(sub)
    print "c%d       = " % (i+1), num.__hex__()
    tot = (tot + num) * rval
  tot = tot % mod1305
  print "r        = ", hexify(r)
  print "m(r)...  = ", tot.__hex__()
  print "k        = ", hexify(k)
  print "n        = ", hexify(n)
  enc = aes_encrypt(k, n)
  print "aes(k,n) = ", hexify(enc)
  enc = str2num_littleend(enc)
  result = (tot + enc) % (1L << 128)
  # Convert to a 16-byte string, little-endian order.
  result = ''.join(map(lambda (i): chr(0xff & (result >> 8*i)), range(16)))
  print "poly1035 = ", hexify(result)
  return result

testvec = [
    { "k" : "\xec\x07\x4c\x83\x55\x80\x74\x17\x01\x42\x5b\x62\x32\x35\xad\xd6",
      "m" : "\xf3\xf6",
      "r" : "\x85\x1f\xc4\x0c\x34\x67\xac\x0b\xe0\x5c\xc2\x04\x04\xf3\xf7\x00",
      "n" : "\xfb\x44\x73\x50\xc4\xe8\x68\xc5\x2a\xc3\x27\x5c\xf9\xd4\x32\x7e",
      "x" : "\xf4\xc6\x33\xc3\x04\x4f\xc1\x45\xf8\x4f\x33\x5c\xb8\x19\x53\xde"
    },
    { "k" : "\x75\xde\xaa\x25\xc0\x9f\x20\x8e\x1d\xc4\xce\x6b\x5c\xad\x3f\xbf",
      "m" : "",
      "r" : "\xa0\xf3\x08\x00\x00\xf4\x64\x00\xd0\xc7\xe9\x07\x6c\x83\x44\x03",
      "n" : "\x61\xee\x09\x21\x8d\x29\xb0\xaa\xed\x7e\x15\x4a\x2c\x55\x09\xcc",
      "x" : "\xdd\x3f\xab\x22\x51\xf1\x1a\xc7\x59\xf0\x88\x71\x29\xcc\x2e\xe7"
    },
    { "k" : "\x6a\xcb\x5f\x61\xa7\x17\x6d\xd3\x20\xc5\xc1\xeb\x2e\xdc\xdc\x74",
      "m" : "\x66\x3c\xea\x19\x0f\xfb\x83\xd8\x95\x93\xf3\xf4\x76\xb6\xbc\x24"
          + "\xd7\xe6\x79\x10\x7e\xa2\x6a\xdb\x8c\xaf\x66\x52\xd0\x65\x61\x36",
      "r" : "\x48\x44\x3d\x0b\xb0\xd2\x11\x09\xc8\x9a\x10\x0b\x5c\xe2\xc2\x08",
      "n" : "\xae\x21\x2a\x55\x39\x97\x29\x59\x5d\xea\x45\x8b\xc6\x21\xff\x0e",
      "x" : "\x0e\xe1\xc1\x6b\xb7\x3f\x0f\x4f\xd1\x98\x81\x75\x3c\x01\xcd\xbe"
    },
    { "k" : "\xe1\xa5\x66\x8a\x4d\x5b\x66\xa5\xf6\x8c\xc5\x42\x4e\xd5\x98\x2d",
      "m" : "\xab\x08\x12\x72\x4a\x7f\x1e\x34\x27\x42\xcb\xed\x37\x4d\x94\xd1"
	  + "\x36\xc6\xb8\x79\x5d\x45\xb3\x81\x98\x30\xf2\xc0\x44\x91\xfa\xf0"
	  + "\x99\x0c\x62\xe4\x8b\x80\x18\xb2\xc3\xe4\xa0\xfa\x31\x34\xcb\x67"
	  + "\xfa\x83\xe1\x58\xc9\x94\xd9\x61\xc4\xcb\x21\x09\x5c\x1b\xf9",
      "r" : "\x12\x97\x6a\x08\xc4\x42\x6d\x0c\xe8\xa8\x24\x07\xc4\xf4\x82\x07",
      "n" : "\x9a\xe8\x31\xe7\x43\x97\x8d\x3a\x23\x52\x7c\x71\x28\x14\x9e\x3a",
      "x" : "\x51\x54\xad\x0d\x2c\xb2\x6e\x01\x27\x4f\xc5\x11\x48\x49\x1f\x1b"
    } ];

def runtests():
  for d in testvec:
    res = poly1305aes(d["k"], d["r"], d["n"], d["m"])
    if res != d["x"]:
      print "COMPUTED RESULT DOESN'T MATCH EXPECTED OUTPUT!"
      print "expected = ", hexify(d["x"])
    print ""