# (c) 04/11/2016, Rolf Rolles, Mobius Strip Reverse Engineering
import z3
import struct

# XOR key for src.bin
KEY_SECTOR = 0x37

# Counter position, as two words
CNTLO      = 0
CNTHI      = 0

# Symbolic names for array positions
POS_CONST0 = 0
POS_KEY0   = 1
POS_KEY2   = 2
POS_KEY4   = 3
POS_KEY6   = 4
POS_CONST2 = 5
POS_NONCE0 = 6
POS_NONCE2 = 7
POS_CNTLO  = 8
POS_CNTHI  = 9
POS_CONST4 = 10
POS_KEY8   = 11
POS_KEY10  = 12
POS_KEY12  = 13
POS_KEY14  = 14
POS_CONST6 = 15

# Constant values and their positions
constvals = [ 0x7865, 0x646E, 0x2D32, 0x6574 ]
constpos  = [ POS_CONST0, POS_CONST2, POS_CONST4, POS_CONST6 ]

# Key byte and word variables
key_bytes = [ z3.BitVec('kb%d' % i,8) for i in xrange(8) ]
key_words = [ None ]*len(key_bytes)

# Representation of key bytes as transformed words
for i in xrange(len(key_bytes)):
    hi = z3.ZeroExt(8,key_bytes[i]) << 9
    lo = z3.ZeroExt(8,key_bytes[i]) + ord('z')
    key_words[i] = hi | lo

# Position of key words within array
keypos  = [ POS_KEY0, POS_KEY2, POS_KEY4, POS_KEY6, POS_KEY8, POS_KEY10, POS_KEY12, POS_KEY14 ]

# Initialize the modified Salsa.16-10 array (word-level values)
def init_array(nonce0,nonce2,cntlo,cnthi):
    initial_array = [None]*16
    for i in xrange(len(constvals)):
        initial_array[constpos[i]] = z3.BitVecVal(constvals[i],16)

    for i in xrange(len(key_words)):
        initial_array[  keypos[i]] = key_words[i]

    initial_array[POS_NONCE0] = z3.BitVecVal(nonce0,16)
    initial_array[POS_NONCE2] = z3.BitVecVal(nonce2,16)
    initial_array[POS_CNTLO]  = z3.BitVecVal(cntlo, 16)
    initial_array[POS_CNTHI]  = z3.BitVecVal(cnthi, 16)
    return initial_array

# Input values to the Salsa round primitive
salsa_steps = [
  (4,0,12,7),
  (8,4,0,9),
  (12,8,4,13),
  (0,12,8,18),

  (9,5,1,7),
  (13,9,5,9),
  (1,13,9,13),
  (5,1,13,18),

  (14,10,6,7),
  (2,14,10,9),
  (6,2,14,13),
  (10,6,2,18),

  (3,15,11,7),
  (7,3,15,9),
  (11,7,3,13),
  (15,11,7,18),

  (1,0,3,7),
  (2,1,0,9),
  (3,2,1,13),
  (0,3,2,18),

  (6,5,4,7),
  (7,6,5,9),
  (4,7,6,13),
  (5,4,7,18),

  (11,10,9,7),
  (8,11,10,9),
  (9,8,11,13),
  (10,9,8,18),

  (12,15,14,7),
  (13,12,15,9),
  (14,13,12,13),
  (15,14,13,18),
]

# z3 representation of the Salsa.16 round primitive
def step(arr,out,inl,inr,rotamt):
    sum = z3.ZeroExt(16,arr[inl] + arr[inr])
    rot = z3.RotateLeft(sum,rotamt)
    arr[out] ^= z3.Extract(15,0,rot)

# Perform one Salsa.16 round
def salsarnd(arr):
    for rnd in salsa_steps:
        step(arr,*rnd)

# Perform all 10 Salsa.16 rounds
def salsa10(arr):
    for i in xrange(10):
        salsarnd(arr)

# Create solver object
s = z3.SolverFor('QF_BV')

# Extract nonce words from the binary nonce data; create initial array
with open("salsa10-nonce.bin", "rb") as f_nonce:
    (n0,n2,n4,n6) = struct.unpack("HHHH", f_nonce.read())
    initial = init_array(n0,n4,CNTLO,CNTHI)
    init_copy = [ x for x in initial ]

# Extract ciphertext words from the binary sector data
with open("salsa10-src.bin", "rb") as f_src:
    srcbytes   = bytearray(f_src.read())
    z3srcwords = [ None ] * len(initial)
    for i in xrange(len(initial)):
        lo = (srcbytes[4*i]   ^ KEY_SECTOR)
        hi = (srcbytes[4*i+1] ^ KEY_SECTOR) << 8
        z3srcwords[i] = z3.BitVecVal(hi | lo,16)

# Specify constraints on key bytes
for k in key_bytes:
    s.add(k != ord('O'))
    s.add(k != ord('I'))
    s.add(k != ord('l'))
    s.add(z3.Or(z3.And(k >= 0x31,k <= 0x39),z3.And(k >= 0x61,k <= 0x78),z3.And(k >= 0x41,k <= 0x58)))

# Perform the Salsa.16 operation
salsa10(initial)

# Specify constraints on outputs
for i in xrange(len(initial)):
    s.add((init_copy[i] + initial[i]) == z3srcwords[i])

# Check SAT, print solution if so
if s.check() == z3.sat:
    m = s.model()
    keybytes = map(lambda k: chr(m[k].as_long()), key_bytes)
    print 'x'.join(keybytes)

else:
    print "UNSAT!"