microcorruption ctf: 'invert' hash for Hollywood level using z3 constriant solver

18_smt2.py

#!/usr/bin/python
from __future__ import division, print_function, unicode_literals
from z3 import *
import binascii, struct

def byteswp(a):
    '''
    Build expression to swap bytes in 16-bit word.
    '''
    return RotateLeft(a,8) # can also be RotateRight, or even Concat(Extract(), Extract())

def hollywood_hash(words):
    '''
    Build expression for hollywood hash.
    '''
    r4 = 0
    r6 = 0
    for w in words:
        #e000: 2455           add @r5, r4
        r4 = r4 + w
        #e000: 8410           swpb r4
        r4 = byteswp(r4)
        #e000: 36e5           xor @r5+, r6
        r6 = r6 ^ w
        #e000: 06e4           xor r4, r6
        #e000: 04e6           xor r6, r4
        #e000: 06e4           xor r4, r6
        (r4,r6) = (r6,r4)

    return (r4,r6)

def main():
    target = (0xfeb1,0x9298) # target values from reverse-engineered code

    for length in range(1,10):
        print('Trying for length %i' % length)

        # Create 16-bit bit vector constants
        words = [BitVec(chr(65+i), 16) for i in range((length+1)//2)]
        # Create hash function for that size
        (r4_f, r6_f) = hollywood_hash(words)

        print(' r4 %s' % r4_f.sexpr())
        print(' r6 %s' % r6_f.sexpr())

        # Create a solver, add constraints
        s = Solver()
        s.add(r4_f == target[0])
        s.add(r6_f == target[1])
        if length & 1:
            s.add(ULT(words[-1], 0x100)) # <0x100

        if s.check() == sat: # satisfyable: print answer and exit
            m = s.model()
            # take values from model for every word
            values = [m[x].as_long() for x in words]
            sol = b''.join(struct.pack('<H', x) for x in values)
            # strip trailing zeros
            while sol[-1] == b'\x00':
                sol = sol[0:-1]
            print('Found solution: \x1b[1;33m%s\x1b[0m' % binascii.b2a_hex(sol))
            exit(0)
        print()

    exit(1)

if __name__ == '__main__':
    main()