import numpy as np
from typing import Dict, List, Iterable
from collections import Counter
from itertools import *
from collections import namedtuple
from Crypto.Cipher import ARC4
import math

def entropy(data: bytes) -> float:
    count = [0] * 256
    dataSize = len(data)
    for b in data:
        count[b] = count[b] + 1
    entropy = 0
    for b in range(256):
        if count[b] / dataSize > 0:
            entropy = entropy + (count[b] / dataSize) * math.log(count[b] / dataSize, 2)
    return -entropy


def codesToStr(codes: Iterable[int]) -> str:
    return ''.join((chr(c) for c in codes))


def bruteforceARC4(cryptogram: bytes, maxKeyLength: int = 5, keyAlphabet=range(256)) -> None:
    entropyThreshold = 6

    for length in range(1, maxKeyLength + 1):
        possibleKeys = product(keyAlphabet, repeat=length)
        for k in possibleKeys:
            k = codesToStr(k)
            cipher = ARC4.new(k)
            recovered = cipher.decrypt(cryptogram)
            e = entropy(recovered)

            if e < entropyThreshold:
                print('\nFound solution!\n')
                print(f'Key: {k}')
                print(f'Notification: {recovered}')
                break

        print(f'No solution for key length {length}')
        print(e)

with open('crypto.rc4', 'rb') as file:
    data = file.read()

    keyAlphabet = range(ord('a'), ord('z') + 1)
    bruteforceARC4(data, 3, keyAlphabet)