import rijndael
import struct
"""
An unoptimized, straightforward reference implementation of VMAC
By Ted Krovetz (tdk@acm.org) -- Modified 22 April 2007
Code assumes (and may break otherwise):
- All bitlengths are a multiple of 8 (ie, bytestrings)
- BLOCKLEN/taglen = 2^i for some integer i.
- L1KEYLEN/64 is an even integer 
"""

# All lengths are bitlengths
BLOCKLEN = 128   # block-length of block cipher in use
KEYLEN   = 128   # block-cipher key-length
L1KEYLEN = 1024  # bits used for l1_hash (compression)

M64  = 0xFFFFFFFFFFFFFFFFL
M126 = 0x3FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFL
MP   = 0x1FFFFFFF1FFFFFFF1FFFFFFF1FFFFFFFL  # Mask for creating poly key
P64  = 0xFFFFFFFFFFFFFEFFL                  # 2^64 - 257
PP   = 0xFFFFFFFF00000000L                  # 2^64 - 2^32
P127 = 0x7FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFL  # 2^127 - 1

def kdf(cipher, index, numbits): # Allows upto 4KB output only
	l = []                 # We will build list of blocks, then join
	n = (numbits+BLOCKLEN-1)//BLOCKLEN  # ceil(numbits/BLOCKLEN)
	prefix = chr(index)+(chr(0) * 14)
	for i in range(n):
		l.append(cipher.encrypt(prefix+chr(i)))
	return (''.join(l))[:(numbits//8)]

def pdf(cipher, nonce, taglen):
	tagsperblock = BLOCKLEN//taglen
	mask = tagsperblock-1  # Assumes tagsperblock in {1,2,4,8,16,...}
	index = ord(nonce[-1]) & mask
	tmpnonce = (chr(0)*(BLOCKLEN//8-len(nonce)))+nonce[:-1]+chr(ord(nonce[-1])-index)
	tmpblock = cipher.encrypt(tmpnonce)
	return tmpblock[index*(taglen//8):(index+1)*(taglen//8)]

def nh(k, m):
	y = 0;
	for i in range(0,len(m),2):
		y += ( ((m[i]+k[i])&M64) * ((m[i+1]+k[i+1])&M64) )
	return y & M126  # Return y mod 2**126

def l1_hash(cipher, m, iter):
	bytesperl1block = L1KEYLEN//8        # Max number of bytes per NH hash
	tmpk = kdf(cipher, 128, L1KEYLEN+128*iter)    # Generate nh key
	fmt = '>%sQ' % str(L1KEYLEN//64)     # Will unpack key block into 64-bit BE words
	k = struct.unpack(fmt,tmpk[-bytesperl1block:])      # Key for nh hash
	t = (len(m)*8+L1KEYLEN-1)//L1KEYLEN  # t = ceil(bitlen(m)/L1KEYLEN)
	y = []                               # Build up result as list of NH hashes
	for i in range(t):                   # Hash each block
		curpos = i * bytesperl1block                             # points to next hash block
		slen = min(bytesperl1block, len(m)-curpos)               # max len, except last block
		padbytesneeded = (16-((slen)%16))%16                     # pad to neaest 16-bytes boundary
		tmpstr = m[curpos:curpos+slen]+chr(0)*padbytesneeded     # Extract/pad next block
		fmt = '<%sQ' % str(len(tmpstr)//8)                       # Unpack into 64-bit LE words
		hstr = struct.unpack(fmt,tmpstr)
		y.append(nh(k,hstr))
	return y                             # Returns empty list if t == 0

def l2_hash(cipher, m, bitlen, iter):
	t = kdf(cipher, 192, 128 * (iter+1))
	t = struct.unpack('>2Q',t[-16:])  # Generate key for poly hash
	k = (long(t[0] & MP) << 64) | (t[1] & MP)       # Construct 128 bits key
	if len(m) == 0:
		y = k
	else:
		y = 1                         # Initialize y for poly eval
		for x in m:                   # Compute polynomial using Horner's rule
			y = (y * k + x) % P127
	return (y + ((bitlen % L1KEYLEN) << 64)) % P127

def l3_hash(cipher, m, iter):
	# Generate key. 1/2^55 chance that it fails the first time (and so loop)
	i = 1
	need = iter + 1
	while need > 0:
		t = kdf(cipher, 224, 128 * i)
		k = struct.unpack('>2Q',t[-16:])  # Generate key for ip hash
		k0 = k[0]
		k1 = k[1]
		i += 1
		if (k0 < P64) and (k1 < P64):
			need -= 1
	m0 = m // PP
	m1 = m % PP
	return ((k0 + m0) * (k1 + m1)) % P64

def vhash(cipher, m, taglen):
	y = []
	for i in range(taglen//64):
		a = l1_hash(cipher,m,i)
		b = l2_hash(cipher,a,len(m)*8,i)
		# print '%X' % b
		y.append(l3_hash(cipher,b,i))
	return y

def vmac(k, m, nonce, taglen):
	if isinstance(k,str): cipher = rijndael.rijndael(k) # if string passed, treat as key
	else: cipher = k                                    # otherwise, treat as cipher
	hashedmessage = vhash(cipher,m,taglen)
	pad = pdf(cipher,nonce,taglen)
	fmt = '>%sQ' % str(taglen//64)
	padnum = struct.unpack(fmt,pad)      # Covert string to number
	tag = []
	for i in range(taglen//64):
		tag.append((padnum[i] + hashedmessage[i]) & M64)
	return tag

# VMAC Test Vectors
def tohex(s):
	res = []
	for c in s:
		res.append('%02X' % ord(c))
	return ''.join(res)

def tags(tag):
	s = []
	for x in tag:
		s.append('%016X' % x)
	return ''.join(s)

vectorkeys = [('abcdefghijklmnop','bcdefghi')]
for k,n in vectorkeys:
	print 'Key   = \'%s\' (%s)' % (k,tohex(k))
	print 'Nonce = \'%s\' %s(%s)' % (n,' '*(len(k)-len(n)),tohex(n))
	print '<<<<<<<<<<<<<<<<<<<<<<<< VMAC Test Vectors >>>>>>>>>>>>>>>>>>>>>>>>'
	print 'Message             64-bit Tag               128-bit Tag'
	print '-------             ----------               -----------'
	for i in [0,1,16,100,10**6]:
		m = 'abc'*i
		print '\'abc\' * %-7d  %s  %s' % (i, tags(vmac(k,m,n,64)), tags(vmac(k,m,n,128)))