-
-
Save graingert/1147493 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import random | |
import copy | |
ALPHABET = ["A","B","C","D","E","F","G","H","I","J","K","L","M","N","O","P","Q","R","S","T","U","V","W","X","Y","Z", " "] | |
class Scrambler: | |
"""A single scrambler unit""" | |
def __init__(self, seed): | |
self.routes = range(len(ALPHABET)) | |
random.seed(seed) | |
random.shuffle(self.routes) | |
def __str__(self): | |
s = '' | |
for i in range(len(ALPHABET)): | |
s += '%s %i => %i\n' % (ALPHABET[i], i, self.routes[i]) | |
return s | |
def forward(self, index): | |
return self.routes[index] | |
def backwards(self, index): | |
return self.routes.index(index) | |
class EnigmaMachine: | |
"""An enigma machine that uses some Scramblers""" | |
def __init__(self): | |
self.scramblers = [] | |
self.offsets = [] | |
self.alphabet = ALPHABET | |
def addScrambler(self, scrambler, offset): | |
self.scramblers.append(scrambler) | |
self.offsets.append(offset) | |
def setPlugboardReplacement(self, a, b): | |
self.alphabet[self.alphabet.index(a)] = b | |
self.alphabet[self.alphabet.index(b)] = a | |
def isReady(self): | |
return len(self.scramblers) >= 1 | |
def reset(self): | |
self.scramblers = [] | |
self.offsets = [] | |
self.alphabet = ALPHABET | |
def scramble(self, char): | |
# For each scrambler, apply offset and trace through | |
char = self.alphabet.index(char.upper()) | |
for i in range(len(self.scramblers)): | |
char = (char + self.offsets[i]) % len(ALPHABET) | |
char = self.scramblers[i].forward(char) | |
char = (len(ALPHABET) - 1) - char | |
for i in reversed(range(len(self.scramblers))): | |
char = self.scramblers[i].backwards(char) | |
char = (char - self.offsets[i]) % len(ALPHABET) | |
# Tick all scramblers on if required | |
for i in range(len(self.offsets)): | |
if i == 0: | |
self.offsets[i] += 1 | |
else: | |
if self.offsets[i - 1] == 0: | |
self.offsets[i] += 1 | |
return self.alphabet[char] | |
def test(): | |
result = "HELLO => " | |
e = EnigmaMachine() | |
e.addScrambler(Scrambler(1), 10) | |
e.addScrambler(Scrambler(2), 100) | |
e.addScrambler(Scrambler(3), 1000) | |
e.setPlugboardReplacement("H", "E") | |
for s in "HELLO WORLD": | |
result += e.scramble(s) | |
result += " => " | |
e = EnigmaMachine() | |
e.addScrambler(Scrambler(1), 10) | |
e.addScrambler(Scrambler(2), 100) | |
e.addScrambler(Scrambler(3), 1000) | |
e.setPlugboardReplacement("H", "E") | |
for s in "EIBDUCKTPCU": | |
result += e.scramble(s) | |
print result | |
if __name__ == '__main__': | |
e = EnigmaMachine() | |
while (True): | |
print "+> ", | |
command = str(raw_input()).upper() | |
if command.startswith("NEW") or command.startswith("RESET"): | |
e.reset() | |
print " # Machine configuration reset" | |
elif command.startswith("ADD"): | |
params = command.replace("ADD", "").replace(" ", "").split(":") | |
e.addScrambler(Scrambler(int(params[0])), (ALPHABET.index(params[1]) % len(ALPHABET))) | |
elif command.startswith("SWAP"): | |
params = command.replace("SWAP", "").replace(" ","").split(":") | |
e.setPlugboardReplacement(params[0], params[1]) | |
elif command.startswith(":"): | |
if e.isReady(): | |
plaintext = command.replace(":", "") | |
cyphertext = '' | |
for s in plaintext: | |
if s in ALPHABET: | |
cyphertext += e.scramble(s) | |
print cyphertext | |
else: | |
print " # Machine not configured!" | |
elif command.startswith("HELP"): | |
print """ # Commands: | |
# new => resets the machine configuration | |
# reset => resets the machine configuration | |
# add <a>:<b> => adds the scrambler with id <id> and oriented to face the letter <orientation | |
# swap <a>:<b> => swaps the letters <a> and <b> on the virtual plugboard | |
# :<plaintext> => scrambles the given plaintext | |
# this will also decrypt a given cyphertext as long as the settings are the same | |
# | |
# test => runs a test encrypt/decrypt and outputs the result | |
# help => prints this message | |
# exit => exits the script""" | |
elif command.startswith("TEST"): | |
test() | |
elif command.startswith("EXIT"): | |
exit() | |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment