from Crypto.Cipher import AES
import base64
import csv

BACK = 'back'
SHIFTR = 'sr'
SHIFTL = 'sl'

def is51234(a, b):
    i = 0
    ai = 0
    bi = 1
    while i < 5:
        if a[ai] != b[bi]:
            return False
        ai = (ai + 1) % 5
        bi = (bi + 1) % 5
        i = i + 1
    return True

def is23451(a, b):
    i = 0
    ai = 1
    bi = 0
    while i < 5:
        if a[ai] != b[bi]:
            return False
        ai = (ai + 1) % 5
        bi = (bi + 1) % 5
        i = i + 1
    return True

def is54321(a, b):
    return a[::-1] == b

def oneDifferent(a, b):
    i = 0
    oneDif = False
    pos = -1
    mode = 0
    while i < 5:
        if a[i] != b[i]:
            if oneDif:
                return False
            else:
                oneDif = True
                pos = i
                if a[i] == 'Circle' and b[i] == 'Square':
                    mode = 2 #counterclock
                elif a[i] == 'Square' and b[i] == 'Diamond':
                    mode = 2
                elif a[i] == 'Diamond' and b[i] == 'Triangle':
                    mode = 2
                elif a[i] == 'Triangle' and b[i] == 'Circle':
                    mode = 2
                elif b[i] == 'Circle' and a[i] == 'Square':
                    mode = 1 #clock
                elif b[i] == 'Square' and a[i] == 'Diamond':
                    mode = 1
                elif b[i] == 'Diamond' and a[i] == 'Triangle':
                    mode = 1
                elif b[i] == 'Triangle' and a[i] == 'Circle':
                    mode = 1
                else:
                   print("FAILURE")

        i = i + 1
    return (pos, mode)

def getProtocol(a, b):
    if is54321(a, b):
        return BACK
    if is23451(a, b):
        return SHIFTL
    if is51234(a, b):
        return SHIFTR
    x = oneDifferent(a, b)
    if x:
        return x
    print("FAILURE2:" + repr(a) + "," + repr(b))
    return False

clock = {'Circle':'Triangle','Triangle':'Diamond','Diamond':'Square','Square':'Circle'}
cclock = {'Circle':'Square','Square':'Diamond','Diamond':'Triangle','Triangle':'Circle'}
def applyClockMode(c, mode):
    if mode == 1:
        return clock[c]
    return cclock[c]

def applyProtocol(a, proto):
    if proto == BACK:
        return a[::-1]
    if proto == SHIFTR:
        return a[4:]+a[0:4]
    if proto == SHIFTL:
        return a[1:5]+a[0:1]
    (pos, mode) = proto
    (A, b, c, d, e) = a
    if pos == 0:
        return (applyClockMode(A, mode), b, c, d, e)
    if pos == 1:
        return (A, applyClockMode(b, mode), c, d, e)
    if pos == 2:
        return (A, b, applyClockMode(c, mode), d, e)
    if pos == 3:
        return (A, b, c, applyClockMode(d, mode), e)
    if pos == 4:
        return (A, b, c, d, applyClockMode(e, mode))
    return False


def extend8(keylist, l):
    keys = []
    i = 0
    while i < l:
        keys.append(keylist[i])
        i = i + 1

    i = l
    j = 0
    while i < 8:
        keys.append(keylist[j])
        i = i + 1
        j = (j + 1) % l
    return keys

out = []
with open("myfile.csv", 'r') as csvfile:
    myreader = csv.reader(csvfile, delimiter=',', quotechar='"')
    for row in myreader:
        out.append(row)

out = out[1:]
mydict = {}
for elt in out:
    key = tuple(elt[2:7])
    check = ''.join(x[0] for x in key)
    if check.lower() != elt[7].lower():
#        print('bad!')
        continue
    if key in mydict:
        new = True
        for test in mydict[key]:
            if ''.join(test[13:17]).lower() == ''.join(elt[13:17]).lower():
                new = False
        if new:
            mydict[key].append(elt)
    else:
        mydict[key] = [elt]

bigcounter = 0
imgs = []

def image(keylist, row):
    global imgs
    global bigcounter
#    print('win')
    try:
        encrypted = row[16]
        IV = row[15]
        key = b''.join([base64.b64decode(x) for x in keylist])
        aes = AES.new(key, AES.MODE_CBC, base64.b64decode(IV))
        z = aes.decrypt(base64.b64decode(encrypted))
#        print('code start:', base64.b64decode(encrypted)[:20])
#        print('iv:', IV)
#        print('decode start:', z[:50])
        if z not in imgs and z.startswith(b'\xff\xd8'):
            print('JPEG found!')
            print('key:', keylist)
            with open('outputfile_{0}.JPEG'.format(bigcounter), 'wb') as g:
                g.write(z)
            bigcounter += 1
            imgs.append(z)
        # if z not in imgs and z.startswith(b'\x00\x00\x00\x18\x66\x74\x79'):
            # print('MPEG found!')
            # print('key:', keylist)
            # with open('outputfile_{0}.mp4'.format(bigcounter), 'wb') as g:
                # g.write(z)
            # bigcounter += 1
            # imgs.append(z)
        # elif z not in imgs:
            # with open('o/{0}'.format(bigcounter), 'wb') as g:
                # g.write(z)
            # bigcounter += 1
            # imgs.append(z)

    except Exception as e:
        neverusing = 0
#        print('bad!!')
#        print(e)
#    print('==============================================')

def dive(key, keylist, depth, startrow, startkey):
    if key in mydict.keys():
        for mutation in mydict[key]:
            newkey = tuple(mutation[8:13])
            proto = getProtocol(key, newkey)
            if not proto:
                continue
            newkey2 = applyProtocol(startkey, proto)

            keylist.append(mutation[14])
            #newkeys = extend8(keylist, depth)
            #image(newkeys, startrow)
            if depth == 8:
                image(keylist, startrow)
                keylist.pop()
                continue
            dive(newkey2, keylist, depth+1, startrow, startkey)
            keylist.pop()

for startkey in mydict.keys():
    mutations = mydict[startkey]
    for mutation in mutations:
        newkey = tuple(mutation[8:13])
        dive(newkey, [mutation[14]], 2, mutation, startkey)