exploit.py

#!/usr/bin/env python3

import re

from z3 import *
from string import printable


def build_alphabet(pattern):
    regex = re.compile(pattern)
    return ''.join(x for x in printable if regex.match(x))


def generate_bitvecs(name, length, size=8):
    return [BitVec(name + str(i), size) for i in range(length)]


def make_solution(ciphertext, pt_length, key_length, constraint):
    pt_vars = generate_bitvecs('pt', pt_length)
    key_vars = generate_bitvecs('key', key_length)
    equations = [pt_vars[0] == constraint]
    for i, ct in enumerate(ciphertext):
        equation = pt_vars[i % pt_length] ^ key_vars[i % key_length] == ct
        equations.append(equation)
    solver = Solver()
    solver.add(equations)
    if solver.check().r == -1:
        return None
    model = solver.model()
    return bytes(model[pt].as_long() for pt in pt_vars)


def brute_constraint(pattern, ciphertext, pt_length, key_length):
    alphabet = build_alphabet(pattern)
    for x in alphabet:
        solution = make_solution(ciphertext, pt_length, key_length, ord(x))
        if solution is not None:
            if re.match(pattern.encode(), solution):
                yield solution


def main():
    pattern = r'^\w+$'
    ciphertext = bytes.fromhex(open('output.txt', 'r').read().strip())
    pt_length = len(ciphertext) // 2
    max_key_length = 32
    for key_length in range(1, max_key_length + 1):
        print('key_length = {0}'.format(key_length))
        for solution in brute_constraint(pattern, ciphertext, pt_length, key_length):
            print('Cup{{{0}}}'.format(solution.decode()))


if __name__ == '__main__':
    main()