#!/usr/bin/env python
def rc4crypt(data, key):
    x = 0
    box = range(256)
    for i in range(256):
        x = (x + box[i] + key[i % len(key)]) % 256
        box[i], box[x] = box[x], box[i]
    x = 0
    y = 0
    out = []
    for byte in data:
        x = (x + 1) % 256
        y = (y + box[x]) % 256
        box[x], box[y] = box[y], box[x]
        out.append(byte ^ box[(box[x] + box[y]) % 256])
    return out

def check_rc4_test_vector():
    data = [0]*16
    key = [0x01,0x02,0x03,0x04,0x05]
    # See RFC 6229
    testVector = [0xb2,0x39,0x63,0x05,0xf0,0x3d,0xc0,0x27,0xcc,0xc3,0x52,0x4a,0x0a,0x11,0x18,0xa8] 

    output = rc4crypt(data, key)
    
    if (output != testVector):
        print "=== Test Vectors Failed ==="
    else:
        print "=== Test Vectors Passed ==="

def get_rand_bytes(numberOfBytes):
    out = []
    file = open("/dev/urandom")
    data = file.read(numberOfBytes)
    for byte in data:
        out.append(ord(byte))
    file.close()
    return out

class Challenger:
    def __init__(self):
        self.whichMessage = self.select_message()
    
    def select_message(self):
        byte = get_rand_bytes(1)[0]
        if byte & 1 == 0:
            return "m0"
        else:
            return "m1"

    def Challenge(self,m0, m1):
        key = get_rand_bytes(16)
        if (self.whichMessage == "m0"):
            return rc4crypt(m0, key)
        elif (self.whichMessage == "m1"):
            return rc4crypt(m1, key)

class Attacker:
    def __init__(self, challenger):
        self.challenger = challenger
    
    def attack(self, numberOfTries):
        byteHistogram = [0] * 256
        m0 = [0,0]
        m1 = [255,255]
        for i in range(1, numberOfTries):
            result = self.challenger.Challenge(m0, m1)
            byteHistogram[result[1]]+=1
        
        if byteHistogram[0] > byteHistogram[255]:
            return "m0"
        else:
            return "m1"

def main():
    import argparse
    parser = argparse.ArgumentParser(description='This program demonstrates a trivial distinguishing attack on RC4 under the same message being encrypted under many keys.')
    parser.add_argument('--trials', type=int,help='Set the number of times the attacker should iterate for.',default=2000)
    
    check_rc4_test_vector()

    challenger = Challenger()
    print "Challenger says: " + challenger.whichMessage    
    
    args = parser.parse_args()
    attacker = Attacker(challenger)
    result = attacker.attack(args.trials)
    print "Attacker says: " + result
    

if __name__ == "__main__":
    main()