from z3 import *

# Eight byte nonce conactenated with 8 null bytes
# Obtained from sector 54
nonce = [0xB0, 0x99, 0x9B, 0x9E, 0xE4, 0xEE, 0x74, 0xC2, 0, 0, 0, 0, 0, 0, 0, 0]


# Verification bytes xored with 0x37
# Obtained from sector 55
# Expanding the key must produce this keystream
targetkeystream = [
    0xC9, 0xF0, 0x00, 0x00, 0x3E, 0xD0, 0x00, 0x00, 0x58, 0xC5, 0x00, 0x00, 0x54, 0x03, 0x00, 0x00,
    0x88, 0xC0, 0x00, 0x00, 0xDC, 0xC8, 0x00, 0x00, 0xE0, 0x71, 0x00, 0x00, 0xC8, 0x85, 0x00, 0x00,
    0x00, 0x04, 0x00, 0x00, 0x00, 0xC1, 0x00, 0x00, 0x65, 0x5A, 0x00, 0x00, 0xCA, 0xAD, 0x00, 0x00,
    0x36, 0xC6, 0x00, 0x00, 0x8E, 0x33, 0x00, 0x00, 0x56, 0x85, 0x00, 0x00, 0xE8, 0xCA, 0x00, 0x00 ]


def rotl(value, shift):
    return (value << shift) |  LShR(value, 32-shift)

def quarterround(y0, y1, y2, y3):
  y1 = y1 ^ rotl(y0 + y3, 7);
  y2 = y2 ^ rotl(y1 + y0, 9);
  y3 = y3 ^ rotl(y2 + y1, 13);
  y0 = y0 ^ rotl(y3 + y2, 18);
  return y0, y1, y2, y3

def rowround(y):
  y[0], y[1], y[2], y[3] = quarterround(y[0], y[1], y[2], y[3])
  y[5], y[6], y[7], y[4] = quarterround(y[5], y[6], y[7], y[4])
  y[10], y[11], y[8], y[9] = quarterround(y[10], y[11], y[8], y[9])
  y[15], y[12], y[13], y[14] = quarterround(y[15], y[12], y[13], y[14])
  return y

def columnround(x):
  x[0], x[4], x[8], x[12] = quarterround(x[0], x[4], x[8], x[12])
  x[5], x[9], x[13], x[1] = quarterround(x[5], x[9], x[13], x[1])
  x[10], x[14], x[2], x[6] = quarterround(x[10], x[14], x[2], x[6])
  x[15], x[3], x[7], x[11] = quarterround(x[15], x[3], x[7], x[11])
  return x

def doubleround(x):
  x = columnround(x)
  x = rowround(x)
  return x

def littleendian(b0, b1):
  return b1, b0

def rev_littleendian(b, w, i):
    b[i + 0] = Extract(7, 0, w)
    b[i + 1] = Extract(15, 8, w)
    b[i + 2] = BitVecVal(0, 8)
    b[i + 3] = BitVecVal(0, 8)
    return b

def hash(seq):
    x = [BitVec('x_%d' %i, 16) for i in range(16)]
    z = [BitVec('z_%d' %i, 16) for i in range(16)]

    for i in range(16):
        b0, b1 = littleendian(seq[4*i], seq[4*i + 1])
        x[i] = z[i] = Concat(BitVecVal(0, 8), b1) + Concat(b0, BitVecVal(0, 8))

    for i in range(10):
        z = doubleround(z)

    for i in range(16):
        z[i] += x[i]
        seq = rev_littleendian(seq, z[i], 4*i)

    return seq

def expand32(k, n, keystream):
    o = [
    [ 'e', 'x', 'p', 'a' ],
    [ 'n', 'd', ' ', '3' ],
    [ '2', '-', 'b', 'y' ],
    [ 't', 'e', ' ', 'k' ]
    ]

    for i in range(0, 64, 20):
        for j in range(4):
            keystream[i+j] = BitVecVal(ord(o[i / 20][j]), 8)

    for i in range(16):
        keystream[4+i]  = k[i]
        keystream[44+i] = k[i+16]
        keystream[24+i] = n[i]

    keystream = hash(keystream)
    return k, n, keystream

def main():
    global nonce, targetkeystream

    n = [BitVecVal(i, 8) for i in nonce]
    keystream = [0] * 64
    key, bvlist = [], []

    for i in range(16):
        bv = BitVec('k_%d' %i, 8)
        bvlist.append(bv)
        key.append(bv+122)
        key.append(bv*2)

    key, n, keystream = expand32(key, n, keystream)
    s = Solver()

    for i in range(len(targetkeystream)):
        s.add(keystream[i] == targetkeystream[i])

    if s.check() == sat:
        print 'Key Found:',
        m = s.model()
        keystring = reduce(lambda acc, x: acc + '1' if m[x] == None else acc + chr(m[x].as_long()), bvlist, '')
        print keystring

if __name__ == '__main__':
    main()