genz3code.py

#!/usr/bin/python2

from hashlib import sha512
import sys

codePath = 'solve.py'

def verify(x, chalbox):
    with open(codePath, mode='w') as f:
        print >> f, 'import sys'
        print >> f, 'from z3 import *'

        length, gates, check = chalbox
        b = [(x >> i) & 1 for i in range(length)]
        for num in range(length):
            print >> f, '{0} = Bool(\'{0}\')'.format(getB(num))
        count = length
        for name, args in gates:
            bName = getB(count)

            if name == 'true':
                b.append(1)
                print >> f, '{0} = True'.format(bName)
            else:
                u1 = b[args[0][0]] ^ args[0][1]
                u1s = getRestrictionStr('Xor', getB(args[0][0]), args[0][1])
                u2 = b[args[1][0]] ^ args[1][1]
                u2s = getRestrictionStr('Xor', getB(args[1][0]), args[1][1])
                if name == 'or':
                    b.append(u1 | u2)
                elif name == 'xor':
                    b.append(u1 ^ u2)
                print >> f, '{0} = {1}'.format(bName, getRestrictionStr(('Or' if name == 'or' else 'Xor'), u1s, u2s))
            count+=1
        
        print >> f, 's = Solver()'
        print >> f, 's.add(Xor({0},{1}) == True)'.format(getB(check[0]), 'True' if check[1] else 'False')
        print >> f, 'r = s.check()'
        print >> f, 'if r == sat:'
        print >> f, '  print(s.model())'
        print >> f, 'else:'
        print >> f, '  print(\'failed\')'

    return b[check[0]] ^ check[1]


def getB(x):
	return 'b{0}'.format(x)

def getRestrictionStr(op,a,b):
	return '{0}({1},{2})'.format(op,a,b)

def dec(x, w):
    z = int(sha512(str(int(x))).hexdigest(), 16)
    return '{:x}'.format(w ^ z).decode('hex')

if __name__ == '__main__':
    if len(sys.argv) < 3:
        print 'Usage: ' + sys.argv[0] + ' <key> <map.txt>'
        print 'Example: Try Running ' + sys.argv[0] + ' 11443513758266689915 map1.txt'
        exit(1)
    with open(sys.argv[2], 'r') as f:
        cipher, chalbox = eval(f.read())
    
    key = int(sys.argv[1]) % (1 << chalbox[0])
    print 'Attempting to decrypt ' + sys.argv[2] + '...'
    if verify(key, chalbox):
        print 'Congrats the flag for ' + sys.argv[2] + ' is:', dec(key, cipher)
    else:
        print 'Wrong key for ' + sys.argv[2] + '.'