from itertools import product
import subprocess

chars = '0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!"#$%&\'()*+,-./:;<=>?@[\\]^_`{|}~ '


def check_pair(script, instr, outstr):
    try:
        result = subprocess.run(
            [
                "/usr/bin/env", "bash", "-c",
                ";\n".join((
                    # Cleanup child processes on exit
                    "trap 'kill -9 $(jobs -p)' SIGINT SIGTERM EXIT",
                    script,
                ))
            ],
            timeout=5,
            input=instr.encode(),
            capture_output=True,
        )
        return result.returncode == 0 and result.stdout == outstr.encode()
    except subprocess.TimeoutExpired:
        return False


def check_pairs(script, pairs):
    for pair in pairs:
        if not check_pair(script, *pair):
            return False
    return True


def find_script(pairs, max_length):
    candidates = []
    chars_count = len(chars)
    total_options = int((chars_count ** (max_length + 1) - 1) / (chars_count - 1))
    current_option = 0

    for length in range(max_length + 1):
        for letters in product(chars, repeat=length):
            current_option += 1

            if current_option % 1000 == 0:
                print(f"Progress: {current_option}/{total_options}")

            script = "".join(letters)

            if check_pairs(script, pairs):
                print("> Found candidate:", script)
                candidates.append(script)

    print("Candidates:", candidates)


if __name__ == '__main__':
    # Identity
    print("\nSearching for identity")
    find_script((
        ("1", "1"),
        ("2", "2"),
        ("3", "3"),
    ), max_length=3)

    # Successor
    # print("\nSearching for successor")
    # find_script((
    #     ("1", "2"),
    #     ("2", "3"),
    #     ("3", "4"),
    # ), max_length=5)