#!/usr/bin/env python3
import sys
import os
from random import randint
from math import ceil
from multiprocessing import Pool


MAX_KEY_SIZE = 2 ** 16 - 1


class mt19937:
    # Source: https://en.wikipedia.org/wiki/Mersenne_Twister
    w = 32
    n = 624
    m = 397
    r = 31

    a = int("9908B0DF", 16)

    u = 11
    d = int("FFFFFFFF", 16)

    s = 7
    b = int("9D2C5680", 16)

    t = 15
    c = int("EFC60000", 16)

    l = 18

    f = 1812433253

    lower_mask = (1 << r) - 1
    w_mask = 2 ** w - 1  # lowest w bits of X
    upper_mask = w_mask & (~lower_mask)

    def __init__(self, seed=5489):
        self.state = [None] * self.n
        self.state[0] = seed
        self.index = 0
        self.seeded_to = 1

    def extract_number(self):
        self.jit_seed()
        self.jit_twist()

        y = self.state[self.index]
        y ^= (y >> self.u) & self.d
        y ^= (y << self.s) & self.b
        y ^= (y << self.t) & self.c
        y ^= y >> self.l

        self.index += 1
        return self.w_mask & y

    def jit_seed(self):
        max_state_index_needed = (self.index + self.m) % self.n
        if not max_state_index_needed < self.seeded_to:
            for i in range(self.seeded_to, max_state_index_needed + 1):
                self.state[i] = self.w_mask & (
                    self.f * (self.state[i - 1] ^ (self.state[i - 1] >> (self.w - 2)))
                    + i
                )
            self.seeded_to = max_state_index_needed + 1

    def jit_twist(self):
        if self.index >= self.n:
            self.index = 0

        i = self.index
        x = (self.state[i] & self.upper_mask) + (
            self.state[(i + 1) % self.n] & self.lower_mask
        )
        xA = x >> 1
        if (x % 2) != 0:
            xA ^= self.a
        self.state[i] = self.state[(i + self.m) % self.n] ^ xA


def random_bytes(size):
    i = randint(0, 2 ** (8 * size))
    return i.to_bytes(size, sys.byteorder)


def random_crypt(suffix):
    data = random_bytes(randint(32, 96)) + suffix
    return crypt(data, randint(0, MAX_KEY_SIZE))


def crypt(data, key):
    rng = mt19937(key)
    data_length = len(data)
    keystream = bytearray()
    for _ in range(ceil(data_length / 4)):
        keystream.extend(rng.extract_number().to_bytes(4, sys.byteorder))

    return xor(data, keystream[:data_length])


def xor(a, b):
    r = int.from_bytes(a, sys.byteorder) ^ int.from_bytes(b, sys.byteorder)
    return r.to_bytes(len(a), sys.byteorder)


def _find_key(data, suffix, range):
    for key in range:
        plain_text = crypt(data, key)
        if plain_text.endswith(suffix):
            return key
    return None


def _find_key_unpacker(args):
    return _find_key(*args)


def find_key(data, suffix, procs=None):
    if procs is None:
        procs = os.cpu_count()

    arguments = []
    range_size = ceil(MAX_KEY_SIZE / procs)
    i = 0
    for _ in range(ceil(MAX_KEY_SIZE / range_size)):
        x = i + range_size
        if x > MAX_KEY_SIZE:
            x = MAX_KEY_SIZE + 1
        arguments.append((data, suffix, range(i, x)))
        i = x

    with Pool(processes=procs) as pool:
        for result in pool.imap_unordered(_find_key_unpacker, arguments):
            if result is not None:
                return result
    return None


def test_crypt():
    d = random_bytes(randint(100, 200))
    s = randint(0, 2 ** 32 - 1)
    b = crypt(d, s)
    if b == d:
        print("Fail: Nothing Changed!")
        return False

    a = crypt(b, s)
    if a != d:
        print("Fail: Operation wasn't invertable!")
        return False
    return True


def main():
    if test_crypt():
        print("Tests Passed!")
    else:
        print("Tests Failed!")

    suffix = b" the msg's suffix"
    target = random_crypt(suffix)
    key = find_key(target, suffix)
    print(f"Found the key: {key}")
    print(f"The Plain Text: {crypt(target, key)}")


if __name__ == "__main__":
    main()
