Mercurial > public > enigma
changeset 27:dc7f939a2ebf
First cut at a command-line app to encrypt/decrypt.
author | Brian Neal <bgneal@gmail.com> |
---|---|
date | Tue, 29 May 2012 17:12:42 -0500 |
parents | abd21cfb67f4 |
children | 067205259796 |
files | enigma/keyfile.py enigma/machine.py enigma/main.py |
diffstat | 3 files changed, 232 insertions(+), 13 deletions(-) [+] |
line wrap: on
line diff
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/enigma/keyfile.py Tue May 29 17:12:42 2012 -0500 @@ -0,0 +1,79 @@ +# Copyright (C) 2012 by Brian Neal. +# This file is part of Py-Enigma, the Enigma Machine simulation. +# Py-Enigma is released under the MIT License (see License.txt). + +"""Contains a function to read key settings from a file. + +A key file is expected to be formatted as one line per day of the month. Each +line consists of a sequence of space separated columns as follows: + +day number - the first column is the day number (1-31). The lines can be in +any order. + +rotor list - the next 3 or 4 columns should be rotor names. + +ring settings - the next 3 or 4 columns should be ring settings. They can be +in either alphabetic (A-Z) or numeric (0-25) formats. + +plugboard settings - the next 10 columns should be plugboard settings. They +can be in either alphabetic (AB CD EF ...) or numeric (1/2 3/4 ...) formats. + +reflector - the last column must be the reflector name. + +Comment lines have a # character in the first column. Blank lines are ignored. + +Each line must either have exactly 18 or 20 columns to be valid. + +""" +import datetime + + +class KeyFileError(Exception): + pass + + +def get_daily_settings(fp, day=None): + """Read and parse a key file for daily key settings. + + fp - a file-like object + + day - specifies the day number to look for in the file (1-31). If day is + None, the day number from today is used. + + Returns a dictionary of keyword arguments for EnigmaMachine.from_key_sheet. + + """ + if day is None: + day = datetime.date.today().day + + for n, line in enumerate(fp): + line = line.strip() + if line == '' or line[0] == '#': + continue + + cols = line.split() + if len(cols) not in [18, 20]: + raise KeyFileError("invalid column count on line %d" % n) + + rotor_count = 3 if len(cols) == 18 else 4 + + try: + day_num = int(cols[0]) + except ValueError: + raise KeyFileError("invalid day on line %d" % n) + + if day_num != day: + continue + + settings = {} + if rotor_count == 3: + settings['rotors'] = cols[1:4] + settings['ring_settings'] = ' '.join(cols[4:7]) + else: + settings['rotors'] = cols[1:5] + settings['ring_settings'] = ' '.join(cols[5:9]) + + settings['plugboard_settings'] = ' '.join(cols[-11:-1]) + settings['reflector'] = cols[-1] + + return settings
--- a/enigma/machine.py Tue May 29 12:43:17 2012 -0500 +++ b/enigma/machine.py Tue May 29 17:12:42 2012 -0500 @@ -10,6 +10,7 @@ from .rotors.factory import create_rotor, create_reflector from .plugboard import Plugboard +from .keyfile import get_daily_settings class EnigmaError(Exception): @@ -18,6 +19,7 @@ # The Enigma keyboard consists of the 26 letters of the alphabet, uppercase # only: KEYBOARD_CHARS = string.ascii_uppercase +KEYBOARD_SET = set(KEYBOARD_CHARS) class EnigmaMachine: @@ -111,6 +113,20 @@ create_reflector(reflector), Plugboard.from_key_sheet(plugboard_settings)) + @classmethod + def from_key_file(cls, fp, day=None): + """Convenience function to read key parameters from a file. + + fp - a file-like object that contains daily key settings + day - the line labeled with the day number (1-31) will be used for the + settings. If day is None, the day number will be determined from today's + date. + + For more information on the file format, see keyfile.py. + + """ + args = get_daily_settings(fp, day) + return cls.from_key_sheet(**args) def set_display(self, val): """Sets the rotor operator windows to 'val'. @@ -143,7 +159,7 @@ The lamp that is lit by this key press is returned as a string. """ - if key not in KEYBOARD_CHARS: + if key not in KEYBOARD_SET: raise EnigmaError('illegal key press %s' % key) # simulate the mechanical action of the machine @@ -207,15 +223,28 @@ return self.plugboard.signal(pos) - def process_text(self, text): + def process_text(self, text, replace_char='X'): """Run the text through the machine, simulating a key press for each letter in the text. + text - the text to process. Note that the text is converted to upper + case before processing. + + replace_char - if text contains a character not on the keyboard, replace + it with replace_char; if replace_char is None the character is dropped + from the message + """ - # TODO: if there is a character not on the keyboard, perform a - # substitution or skip it. result = [] for key in text: - result.append(self.key_press(key)) + c = key.upper() + + if c not in KEYBOARD_SET: + if replace_char: + c = replace_char + else: + continue # ignore it + + result.append(self.key_press(c)) return ''.join(result)
--- a/enigma/main.py Tue May 29 12:43:17 2012 -0500 +++ b/enigma/main.py Tue May 29 17:12:42 2012 -0500 @@ -2,20 +2,131 @@ # This file is part of Py-Enigma, the Enigma Machine simulation. # Py-Enigma is released under the MIT License (see License.txt). -from .machine import EnigmaMachine +"""Provide an example command-line app that can setup an EnigmaMachine and +process text. + +""" + +import argparse +import sys + +from .machine import EnigmaMachine, EnigmaError +from .rotors import RotorError + + +PROG_DESC = 'Encrypt/decrypt text according to Enigma machine key settings' + +HELP_EPILOG = """\ +Key settings can either be specified by command-line arguments, or read +from a key file. If reading from a key file, the line labeled with the +current day number is used unless the --day argument is provided. + +Text to process can be supplied 3 ways: + + if --text=TEXT is present TEXT is processed + if --file=FILE is present the contents of FILE are processed + otherwise the text is read from standard input + +Examples: + + $ %(prog)s --key-file=enigma.keys -t HELLOXWORLDX + $ %(prog)s -r III IV V -s 0 1 2 -p AB CD EF GH IJ KL MN -u B + $ %(prog)s -r Beta III IV V -s A B C D -p 1/2 3/4 5/6 -u B-Thin + +""" + +def create_from_key_sheet(filename, day=None): + """Create an EnigmaMachine from a daily key sheet.""" + + with open(filename, 'r') as f: + return EnigmaMachine.from_key_file(f, day) + + +def create_from_args(parser, args): + """Create an EnigmaMachine from command-line specs.""" + + if args.rotors is None: + parser.error("Please specify 3 or 4 rotors; e.g. II IV V") + elif len(args.rotors) not in [3, 4]: + parser.error("Expecting 3 or 4 rotors; %d supplied" % len(args.rotors)) + + if args.text and args.file: + parser.error("Please specify --text or --file, but not both") + + ring_settings = ' '.join(args.ring_settings) if args.ring_settings else None + plugboard = ' '.join(args.plugboard) if args.plugboard else None + + return EnigmaMachine.from_key_sheet(rotors=args.rotors, + ring_settings=ring_settings, + plugboard_settings=plugboard, + reflector=args.reflector) def main(): - machine = EnigmaMachine.from_key_sheet( - rotors=['I', 'II', 'III'], - reflector='B') + parser = argparse.ArgumentParser(description=PROG_DESC, epilog=HELP_EPILOG, + formatter_class=argparse.RawDescriptionHelpFormatter) + + parser.add_argument('-k', '--key-sheet', + help='path to key sheet for daily settings') + parser.add_argument('-d', '--day', type=int, default=None, + help='use the settings for day DAY when reading keysheet') + parser.add_argument('-r', '--rotors', nargs='+', metavar='ROTOR', + help='rotor list ordered from left to right; e.g III IV I') + parser.add_argument('-i', '--ring-settings', nargs='+', + metavar='RING_SETTING', + help='ring setting list from left to right; e.g. A A J') + parser.add_argument('-p', '--plugboard', nargs='+', metavar='PLUGBOARD', + help='plugboard settings') + parser.add_argument('-u', '--reflector', help='reflector name') + parser.add_argument('-s', '--start', help='starting position') + parser.add_argument('-t', '--text', help='text to process') + parser.add_argument('-f', '--file', help='input file to process') + parser.add_argument('-x', '--replace-char', default='X', + help=('if the input text contains chars not found on the enigma' + ' keyboard, replace with this char [default: %(default)s]')) + parser.add_argument('-z', '--delete-chars', default=False, + action='store_true', + help=('if the input text contains chars not found on the enigma' + ' keyboard, delete them from the input')) - machine.set_display('AAA') - cipher_text = machine.process_text('AAAAA') + args = parser.parse_args() - print(cipher_text) + if args.key_sheet and (args.rotors or args.ring_settings or args.plugboard + or args.reflector): + parser.error("Please specify either a keysheet or command-line key " + "settings, but not both") + + if args.start is None: + parser.error("Please specify a start position") + + if args.key_sheet: + machine = create_from_key_sheet(args.key_sheet, args.day) + else: + machine = create_from_args(parser, args) + + if args.text: + text = args.text + elif args.file: + with open(args.file, 'r') as f: + text = f.read() + else: + text = input('--> ') + + replace_char = args.replace_char if not args.delete_chars else None + + machine.set_display(args.start) + + s = machine.process_text(text, replace_char=replace_char) + print(s) + + +def console_main(): + try: + main() + except (IOError, EnigmaError, RotorError) as ex: + sys.stderr.write("%s\n" % ex) if __name__ == '__main__': - main() + console_main()