http://site.douban.com/Jobs/widget/notes/277664/note/518314301/

douban_puzzle.py

import itertools


ALPHABET = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
ALPHABET_LEN = len(ALPHABET)
ALPHABET_INDEX = {c: i for i, c in enumerate(ALPHABET)}


def a2n(letter):
    return ALPHABET_INDEX[letter]


def n2a(number):
    return ALPHABET[number]


class Enigma(object):

    def __init__(self, rotors, reflector):
        self.rotors = rotors
        self.reflector = reflector
        self.setup()

    def setup(self):
        for i, r in enumerate(self.rotors):
            if i < len(self.rotors) - 1:
                self.rotors[i].next_rotor = self.rotors[i + 1]

        self.processors = (self.rotors +
                           [self.reflector] +
                           [r.inverse for r in reversed(self.rotors)])

    def __call__(self, text):
        result = ''
        for c in text:
            self.rotors[0].rotate()
            n = a2n(c)
            for f in self.processors:
                n = f(n)
            result += n2a(n)
        return result


class Rotor(object):

    MODELS = [
        ('EKMFLGDQVZNTOWYHXUSPAIBRCJ', 'Q'),  # I
        ('AJDKSIRUXBLHWTMCQGZNPYFVOE', 'E'),  # II
        ('BDFHJLCPRTXVZNYEIWGAKMUSQO', 'V'),  # III
        ('ESOVPZJAYQUIRHXLNFTGKDCMWB', 'J'),  # IV
        ('VZBRGITYUPSDNHLXAWMJQOFECK', 'Z'),  # V
    ]

    def __init__(self, model=0, position=0, next_rotor=None):
        self.wiring_string, self.turnover_letter = self.MODELS[model]
        self.position = position
        self.next_rotor = next_rotor

        self.turnover = a2n(self.turnover_letter)

        self.wiring = range(len(self.wiring_string))
        for i, c in enumerate(self.wiring_string):
            self.wiring[i] = a2n(c)

        self.inverse_wiring = range(len(self.wiring_string))
        for i, c in enumerate(self.wiring):
            self.inverse_wiring[c] = i

    def __call__(self, letter, inverse=False):
        x = (self.position + letter) % ALPHABET_LEN
        x = self.wiring[x] if not inverse else self.inverse_wiring[x]
        x = (ALPHABET_LEN - self.position + x) % ALPHABET_LEN
        return x

    def inverse(self, letter):
        return self(letter, inverse=True)

    def rotate(self):
        if self.position == self.turnover and self.next_rotor:
            self.next_rotor.rotate()
        self.position += 1
        if self.position >= ALPHABET_LEN:
            self.position = 0


class Reflector(object):

    MODELS = [
        'YRUHQSLDPXNGOKMIEBFZCWVJAT',  # B
        'FVPJIAOYEDRZXWGCTKUQSBNMHL',  # C
        'ENKQAUYWJICOPBLMDXZVFTHRGS',  # B thin
        'RDOBJNTKVEHMLFCWZAXGYIPSUQ',  # C thin
    ]

    def __init__(self, model):
        self.wiring_string = self.MODELS[model]
        self.wiring = range(len(self.wiring_string))
        for i, c in enumerate(self.wiring_string):
            self.wiring[i] = a2n(c)

    def __call__(self, letter):
        return self.wiring[letter]


def decrypt(ciphertext, words, num_rotors=3):
    for reflector_model in range(len(Reflector.MODELS)):
        model_gen = itertools.permutations(range(len(Rotor.MODELS)), num_rotors)
        for rotor_models in model_gen:
            pos_gen = itertools.product(range(ALPHABET_LEN), repeat=num_rotors)
            for rotor_positions in pos_gen:
                rotors = [Rotor(model=m, position=p)
                          for m, p in zip(rotor_models, rotor_positions)]
                reflector = Reflector(reflector_model)

                enigma = Enigma(rotors=rotors, reflector=reflector)
                text = enigma(ciphertext)
                if any(w in text for w in words):
                    print text


def main():
    ciphertext = 'QCARJBVMRVJKUKMXXRCO'
    words = ('DOUBAN', )
    decrypt(ciphertext, words)  # DOUBANWANTAWESOMEYOU


if __name__ == '__main__':
    main()