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()