#! /usr/bin/python

import catacomb as CAT

I = 1
A = 2
C = 4
T = 8
M = 16
INIT = 256

def dump_keccak_state(what, state):
  buf = CAT.ReadBuffer(state + CAT.ByteString.zero(200 - len(state)))
  print ";; %s [round 0]" % what
  print ";; uncomplemented state..."
  for i in xrange(5):
    print ";;\t%s" % " ".join(["%016x" % buf.getu64l() for j in xrange(5)])
  print

class Strobe (object):
  def __init__(me, sec = 128, _copy = None):
    if _copy is None:
      me.k = CAT.Keccak1600()
      me.r = (1600 - 2*sec)/8
      me.f = 0
      me.p0 = 0
      buf = CAT.WriteBuffer()
      buf.putu8(1).putu8(me.r)
      buf.putu8(1).putu8(0)
      buf.putu8(1).putu8(12*8).put('STROBEv1.0.2')
      ##dump_keccak_state("init", buf.contents)
      me.k.mix(buf.contents).step()
      me.buf = CAT.WriteBuffer()
      me.mask = me.k.extract(me.r - 2)
      ##dump_keccak_state("final", me.mask)
    else:
      me.k = _copy.k.dup()
      me.r = _copy.r
      me.f = _copy.f
      me.p0 = _copy.p0
      me.buf = CAT.WriteBuffer().put(_copy.buf.contents)
      me.mask = me._copy.mask

  def _crank(me):
    me.buf.putu8(me.p0); me.p0 = 0
    if me.buf.size == me.r - 1: me.buf.putu8(0x84)
    else: me.buf.putu8(0x04).zero(me.r - me.buf.size - 1).putu8(0x80)
    ##dump_keccak_state("buffer", me.buf.contents)
    ##dump_keccak_state("init", me.buf.contents ^ me.k.extract(me.r))
    me.k.mix(me.buf.contents).step()
    me.mask = me.k.extract(me.r - 2)
    ##dump_keccak_state("final", me.mask)
    me.buf = CAT.WriteBuffer()

  def process(me, op, input):
    if not op&T:
      opf = op
    else:
      if not me.f&INIT: me.f |= INIT | (op&I)
      opf = op ^ (me.f&255)
    me.buf.putu8(me.p0); me.p0 = me.buf.size
    if me.buf.size >= me.r - 2: me._crank()
    me.buf.putu8(opf)
    if op&C or me.buf.size >= me.r - 2: me._crank()

    if op&(I | T) == I or op&(I | A) == 0: buf = CAT.ByteString.zero(input)
    else: buf = CAT.ByteString(input)
    out = CAT.WriteBuffer()
    i, j, sz = 0, me.buf.size, len(buf)

    while i < sz:
      spare = me.r - j - 2
      n = min(spare, sz - i); x = buf[i:i + n]; i += n
      if not op&C: out.put(x); me.buf.put(x)
      else:
        y = x ^ me.mask[j:j + n]; out.put(y)
        if op&I or not op&T: me.buf.put(y)
        else: me.buf.put(x)
      if n == spare: me._crank(); j = 0

    if op&(I | A | C | T) != (I | C | T): return out.contents
    elif out.contents == CAT.ByteString.zero(out.size): return me
    else: raise ValueError, 'verify failed'

st0 = Strobe()
st1 = Strobe()

st0.process(A | M, 'testing')
st1.process(A | M, 'testing')

print 'prf(10) -> %s' % hex(st0.process(I | A | C, 10))
st1.process(I | A | C, 10)

print 'ad("Hello")'
st0.process(A, 'Hello')
st1.process(A, 'Hello')

c0 = st0.process(A | C | T, 'World');
st1.process(I | A | C | T, c0)
print 'encout("World") -> %s' %  hex(c0)

print 'clrout("foo")'
st0.process(A | T, 'foo');
st1.process(I | A | T, 'foo')

print 'clrin("bar")'
st1.process(A | T, 'bar')
st0.process(I | A | T, 'bar')

m1 = st0.process(I | A | C | T, 'baz')
st1.process(A | C | T, m1) #  backwards in time
print 'encin("baz") -> %s' % hex(m1)

for i in xrange(200):
  c = st0.process(A | C | T, i*'X')
  m = st1.process(I | A | C | T, c)
  print "  0 encout\n\t=%s\n\t!%s;" % (i*'X', hex(c))
  print "  1 encin\n\t!%s\n\t=%s;" % (hex(c), i*'X')

r0 = st0.process(I | A | C, 123)
r1 = st1.process(I | A | C, 123)
print 'prf(123) -> %s' % hex(r0)
assert r0 == r1

t = st0.process(C | T, 16)
print 'macout(16) -> %s' % hex(t)
st1.process(I | C | T, t)

k = 'this is my super-secret key'
print 'key("%s")' % k
st0.process(A | C, k)
st1.process(A | C, k)

r0 = st0.process(I | A | C, 32)
r1 = st1.process(I | A | C, 32)
print 'prf(32) -> %s' % hex(r0)
assert r0 == r1

print 'ratchet(32)'
st0.process(C, 32)
st1.process(C, 32)

t0 = CAT.bytes('b2084ebdfabd50768c91eebc190132cc')
st0.process(I | C | T, t0)
t1 = st1.process(C | T, 16)
assert t1 == t0
print 'macin(%s)' % hex(t0)