"""Guided REA reproduction of joseph's DownUnderCTF 2023 masked-squares task.

Install z3-solver, then run: python solve.py [checks.json]
The constraints are extracted from the original ELF through REA/Ghidra.
The alphabet and DUCTF{...} shape below are analyst assumptions.
This returns one accepted candidate; it does not prove uniqueness.
"""
import json
import pathlib
import sys
from z3 import Int, Or, Solver, Sum, sat

source = pathlib.Path(sys.argv[1]) if len(sys.argv) > 1 else pathlib.Path(__file__).with_name("checks.json")
data = json.loads(source.read_text())
codes = [Int(f"code_{i}") for i in range(36)]
solver = Solver()
solver.set(timeout=60000)
alphabet = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_{}"
for code in codes:
    solver.add(Or(*[code == ord(char) for char in alphabet]))
for index, char in enumerate("DUCTF{"):
    solver.add(codes[index] == ord(char))
solver.add(codes[-1] == ord("}"))
for check in data["checks"]:
    solver.add(Sum([codes[index] for index in check["selected"]]) == check["target"])
status = solver.check()
if status != sat:
    raise SystemExit(f"No candidate: {status}; {solver.reason_unknown()}")
model = solver.model()
candidate = "".join(chr(model.eval(code).as_long()) for code in codes)
assert all(sum(ord(candidate[i]) for i in c["selected"]) == c["target"] for c in data["checks"])
print(json.dumps({"candidate": candidate, "checksPassed": len(data["checks"]), "z3": __import__("z3").get_version_string(), "uniquenessProven": False}))