#!/usr/bin/env python3
# -*- coding:utf-8 -*-
import random
import sys
import secret

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 randString(n):
    return "".join([chr(random.randint(0, 255)) for _ in range(n)])

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


if __name__ == "__main__":
    if len(sys.argv) != 2:
        print "Usage : %s <file>" % (sys.argv[0])
        sys.exit(-1)

    infile = sys.argv[1]
    outfile = infile + ".enc"
    nonce = randString(8)
    s1, s2 = genSubkeys(nonce)
    data = open(infile, "r").read()
    data = rc4(s1, data)
    data = rc4(s2, data)
    data = nonce + data
    open(outfile, "w").write(data)