Cifra Enigma

# -*- coding: utf-8 -*-
# Uma implementação da cifra Enigma
# Autor: Bruno Flávio de Castro Ribeiro
# A Criptanálise da ENIGMA: 1932-1939
# Tese 217

# Exemplo
# Rotores utilizados: III, II e I (posicao direita, meio e esquerda respectivamente)
# Posicoes: P, J e D
# Aneis ajustados para as posicoes: K, W e D
#
# x = enigma([rotor(rotorIII), rotor(rotorII), rotor(rotorI)], 'KWD', 'PJD', 'WFR', plugboard(painel1), mirror(reflectorB))
# x.egm('MAQUINAENIGMA')

# Componentes da Enigma
import exceptions

class unsoundComponent(Exception):
    pass

class unsoundPlugboard(Exception):
    pass

class Ething(object):
    def verify(self):
        foo = {}
        for i in self.table.keys():
            foo[self.table[i]] = True

        if len(self.table.keys()) != 26 or len(foo.keys()) != 26:
            raise unsoundComponent

    def trad(self, c):
        return self.table[c]

    def tradinv(self, c):
        return self.tableinv[c]

    def tord(self, c):
        return ord(c) - ord('A')

    def mostra(self, c):
        return self.table

class rotor(Ething):
    def __init__(self, perm):
        self.table = {}
        self.tableinv = {}
        foo = 0

        for i in perm:
            self.table[foo] = self.tord(i)
            foo = foo + 1
        self.verify()

        for j in range(len(self.table.keys())):
            self.tableinv[self.table.values()[j]] = self.table.keys()[j]
        self.verify()

class mirror(Ething):
    def __init__(self, perm):
        self.table = {}

        for i in perm:
            self.table[self.tord(i[0])] = self.tord(i[1])
            self.table[self.tord(i[1])] = self.tord(i[0])
        self.verify()

class plugboard(Ething):
    def __init__(self, perm):
        self.table = {}
        for i in range(26):
            self.table[i] = i

        for i in perm:
            self.table[self.tord(i[0])] = self.tord(i[1])
            self.table[self.tord(i[1])] = self.tord(i[0])
        self.verify()
        self.verifyplugboard()

    def verifyplugboard(self):
        foo = {}
        for i in range(len(self.table.keys())):
            if self.table[i] != i:
                foo[i] = True

        if len(foo.keys()) > 20:
            raise unsoundPlugboard

# Plugboards, rotores e reflectores da Enigma
painel0 = []
painel1 = ['UA', 'PF', 'RQ', 'SO', 'NI', 'EY', 'BG', 'HL', 'TX', 'ZJ']
painel2 = ['UB', 'PS', 'RQ', 'FO', 'NY', 'EI', 'AG', 'HL', 'TX', 'ZJ']
painel3 = ['UC', 'PF', 'RQ', 'SO', 'NI', 'EY', 'BG', 'HL', 'TX', 'ZJ']
painel4 = ['UO', 'PF', 'RQ', 'SA', 'NI', 'EY', 'BG', 'HL', 'TX', 'ZJ']
painel5 = ['UP', 'AF', 'RQ', 'SO', 'NI', 'EY', 'BG', 'HL', 'TX', 'ZJ']
painel6 = ['UH', 'PF', 'RQ', 'SO', 'NI', 'EY', 'BG', 'AL', 'TX', 'ZJ']
painel7 = ['UX', 'PF', 'RQ', 'SO', 'NI', 'EY', 'BG', 'HL', 'TA', 'ZJ']
painel8 = ['UI', 'PF', 'RQ', 'SO', 'NA', 'EY', 'BG', 'HL', 'TX', 'ZJ']
painel9 = ['EA', 'PF', 'RQ', 'SO', 'NI', 'UY', 'BG', 'HL', 'TX', 'ZJ']
painel10 = ['UA', 'RF', 'PQ', 'SO', 'NI', 'EY', 'BG', 'HL', 'TX', 'ZJ']

rotorI = 'EKMFLGDQVZNTOWYHXUSPAIBRCJ'
rotorII = 'AJDKSIRUXBLHWTMCQGZNPYFVOE'
rotorIII = 'BDFHJLCPRTXVZNYEIWGAKMUSQO'
rotorIV = 'ESOVPZJAYQUIRHXLNFTGKDCMWB'
rotorV = 'VZBRGITYUPSDNHLXAWMJQOFECK'

reflectorB = ['AY', 'BR', 'CU', 'DH', 'EQ', 'FS', 'GL', 'IP', 'JX', 'KN', 'MO', 'TZ', 'VW']
reflectorC = ['AF', 'BV', 'CP', 'DJ', 'EI', 'GO', 'HY', 'KR', 'LZ', 'MX', 'NW', 'TQ', 'SU']

# Movimento da Enigma
def Rot(c, s, i, e, n, t):
    return (c + s[t] + R(t, n, i, e)) % 26

def Roti(c, s, i, e, n, t):
    return (c - s[t] - R(t, n, i, e)) % 26

def R(t, n, i, e):
    if t == 0:
        return n + 1
    else:
        if t == 1 and i[t] + ((i[t-1] - e[t-1]) % 26 + n) / 26 == e[t] - 1:
            return ((i[t-1] - e[t-1]) % 26 + R(t - 1, n, i, e)) / 26 + 1
        else:
            return ((i[t-1] - e[t-1]) % 26 + R(t - 1, n, i, e)) / 26

class enigma:
    def __init__(self, rotlist, rings, pinit, notch, painel, reflector):
        self.rotlist = rotlist
        self.rings = rings
        self.pinit = pinit
        self.notch = notch
        self.painel = painel
        self.reflector = reflector

    def egm(self, msg):
        msgcod = ''
        lplaint = [0 for i in range(len(msg))]
        pinitial = [0 for i in range(len(self.pinit))]
        rings = [0 for i in range(len(self.rings))]
        shift = [0 for i in range(len(self.rings))]
        notch = [0 for i in range(len(self.notch))]
        codtext = {}

        for n in range(len(msg)):
            lplaint[n] = ord(msg[n]) - ord('A')
            c0 = self.painel.trad(lplaint[n])
            ct = c0

            for j in range(len(self.rotlist)):
                pinitial[j] = ord(self.pinit[j]) - ord('A')
                rings[j] = ord(self.rings[j]) - ord('A')
                shift[j] = (pinitial[j] - rings[j]) % 26
                notch[j] = ord(self.notch[j]) - ord('A')

                c1 = Rot(ct, shift, pinitial, notch, n, j)
                c2 = self.rotlist[j].trad(c1)
                c3 = Roti(c2, shift, pinitial, notch, n, j)
                ct = c3

            cr = self.reflector.trad(ct)

            for k in range(len(self.rotlist)):
                c_1 = Rot(cr, shift, pinitial, notch, n, len(self.rotlist) - 1 - k)
                c_2 = self.rotlist[len(self.rotlist) - 1 - k].tradinv(c_1)
                c_3 = Roti(c_2, shift, pinitial, notch, n, len(self.rotlist) - 1 - k)
                cr = c_3

            cf = self.painel.trad(cr)
            codtext[lplaint[n]] = cf
            msgcod = msgcod + chr(cf + 65)

        return msgcod

# As seis permutacoes diarias
alfa = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'

def state():
    R = []
    H = []

    for k in alfa:
        for s in alfa:
            for t in alfa:
                pinit = k + s + t

                for j in range(6):
                    L = []
                    A = num2let((let2num(pinit[0]) + j) % 26)
                    B = pinit[1]
                    C = pinit[2]
                    pin = A + B + C
                    notch = num2let((let2num(pin[0]) + 7) % 26) + num2let((let2num(pin[1]) + 7) % 26) + 'A'

                    for i in alfa:
                        x = enigma([rotor(rotorII), rotor(rotorIII), rotor(rotorI)], 'AAA', pin, notch, plugboard(painel0), mirror(reflectorB))
                        y = x.egm(i)
                        L.append(i)
                        L.append(y)
                        W = list2dict(L)
                        R.append(W)

                    H.append(R)
                    R = []
    return H

def list2dict(L):
    D = {}
    c = len(L)
    for r in range(c):
        D[L[(2 * r) % c]] = L[(2 * r + 1) % c]

    return D

def let2num(A):
    return ord(A) - 65

def num2let(A):
    return chr((A % 26 + 65))

# Ciclometro
# PS: processo 30m aprox
import pickle

f = open('catalogo.pck', 'w')
alfa = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'

def ciclometer():
    C = []
    for a in range(len(alfa)):
        for b in range(len(alfa)):
            for c in range(len(alfa)):
                pinit1 = alfa[a] + alfa[b] + alfa[c]
                pinit2 = num2let(let2num(pinit1[0]) + 3) + pinit1[1:3]
                A, B, D = [], [], []
                notch = num2let(let2num(pinit1[0]) + 7) + num2let(let2num(pinit1[1]) + 7) + 'A'

                for i in range(len(alfa)):
                    if alfa[i] not in B:
                        let = alfa[i]
                        X = [let]
                        lit = ''
                        while lit != X[0]:
                            x = enigma([rotor(rotorIII), rotor(rotorII), rotor(rotorI)], 'AAA', pinit1, notch, plugboard(painel0), mirror(reflectorB))
                            w = enigma([rotor(rotorIII), rotor(rotorII), rotor(rotorI)], 'AAA', pinit2, notch, plugboard(painel0), mirror(reflectorB))

                        y = x.egm(let)
                        z = w.egm(y)
                        X.append(y)
                        X.append(z)
                        let = z
                        lit = let
                    B = B + X
                    A.append(X[0:len(X) - 1])

                D = D + A
                W = [len(D) * 2]

                for i in range(len(D)):
                    W = W + [len(D[i]) / 2] + [len(D[i]) / 2]
                    C.append(W + [pinit1])

    pickle.dump(C, f)
    f.close()

def let2num(A):
    return ord(A) - 65

def num2let(A):
    return chr((A % 26 + 65))