#!/usr/bin/env python3
# -*- coding:utf-8 -*-
import pickle
import itertools
import zipfile
import string

def rc4(key, text):
    def KSA(key):
        key_length = len(key)
        S = range(256)
        j = 0
        for i in range(256):
            j = (j + S[i] + key[i % key_length]) % 256
            S[i], S[j] = S[j], S[i]  # swap values
        return S

    def PRGA(S):
        i = 0
        j = 0
        while True:
            i = (i + 1) % 256
            j = (j + S[i]) % 256
            S[i], S[j] = S[j], S[i]
            K = S[(S[i] + S[j]) % 256]
            yield K

    def get_keystream(key):
        S = KSA(key)
        return PRGA(S)

    keyBytes = [ord(c) for c in key]
    keystream = get_keystream(keyBytes)
    r = ""
    for c in text:
        r += chr(ord(c) ^ next(keystream))
    return r

def genSubkeys(nonce):
    s1 = ""
    for i in range(3):
        s1 += nonce[3*i:3*(i+1)]+KEY[i]
    s2 = ""
    for i in range(3):
        s2 += nonce[::-1][3*i:3*(i+1)]+KEY[i+3]
    return s1, s2

if __name__ == "__main__":
    # open the encrypted file
    f = open("flag.txt.zip.enc", "r").read()
    nonce = f[:8]
    data = f[8:]
    magic = "PK\x03\x04"

    print("Starting attack, this takes about 2 minutes to complete")
    save = "save.bin"
    # attempt to load saved progress
    table = {}
    print("Attempting to load progress...")
    try:
        table = pickle.load(file(save, "r"))
    except:
        table = {}
    if table == {}:
        print("Generating first table...")
        # generate first table
        for t in itertools.product(string.printable, repeat=3):
            KEY = "".join(t) + "AAA"
            subkey1, _ = genSubkeys(nonce)
            encrypted = rc4(subkey1, magic)
            table[encrypted] = KEY
        # save in a file for later in case a mistake was made so we don't have to recalculate this step
        pickle.dump(table, file(save,"w"))

    # search for matching candidates
    print("Searching for matches... (progress saved)")
    # values = list(table.values())
    # keys = list(table.keys())
    for t in itertools.product(string.printable, repeat=3):
        KEY = "AAA" + "".join(t)
        _, subkey2 = genSubkeys(nonce)
        decrypted = rc4(subkey2, data[:4])
        value = table.get(decrypted)
        if value:
            # index = values.index(decrypted)
            KEY = value[:3] + KEY[3:]
            # match found attempt to decrypt whole file
            s1, s2 = genSubkeys(nonce)
            decryptedZip = data
            decryptedZip = rc4(s1, decryptedZip)
            decryptedZip = rc4(s2, decryptedZip)
            open("flag.zip", "w").write(decryptedZip)
            try:
                zip_ref = zipfile.ZipFile("flag.zip", "r")
                zip_ref.extractall(".")
                zip_ref.close()
                print("Solution found, key = "+repr(KEY))
                print(open("flag.txt", "r").read())
                break
            except:
                pass