# HG changeset patch
# User Brian Neal <bgneal@gmail.com>
# Date 1338251395 18000
# Node ID f0bb15d17556f7927e9c3208d941880da3ede36b
# Parent  ee1c84475eda7da1d48aad11aca0384d444e314c
Wrote a test using the Kriegsmarine example in Dirk Rijmenants'
simulator manual. This revealed a bug in EnigmaMachine.set_display()
for 4 rotor machines. Fixed. This was the first time we have tested a
4 rotor model!

diff -r ee1c84475eda -r f0bb15d17556 enigma/machine.py
--- a/enigma/machine.py	Mon May 28 16:37:22 2012 -0500
+++ b/enigma/machine.py	Mon May 28 19:29:55 2012 -0500
@@ -105,11 +105,11 @@
         window from left to right.
 
         """
-        if len(val) != 3:
-            raise EnigmaError("Bad display value")
+        if len(val) != self.rotor_count:
+            raise EnigmaError("Incorrect length for display value")
 
         for i, rotor in enumerate(reversed(self.rotors)):
-            rotor.set_display(val[2 - i])
+            rotor.set_display(val[-1 - i])
 
     def get_display(self):
         """Returns the operator display as a string."""
diff -r ee1c84475eda -r f0bb15d17556 enigma/tests/test_enigma.py
--- a/enigma/tests/test_enigma.py	Mon May 28 16:37:22 2012 -0500
+++ b/enigma/tests/test_enigma.py	Mon May 28 19:29:55 2012 -0500
@@ -51,8 +51,12 @@
 
 
 class ActualDecryptTestCase(unittest.TestCase):
-    """This example taken from Dirk Rijmenants' simulator manual."""
-
+    """This example taken from Dirk Rijmenants' simulator manual.
+    
+    It is credited to Frode Weierud and Geoff Sullivan.
+        http://cryptocellar.com
+    
+    """
     def setUp(self):
 
         ring_settings = [ord(c) - ord('A') for c in 'BUL']
@@ -132,3 +136,57 @@
         )
 
         self.decrypt('CRS', 'YPJ', ciphertext, truth_data)
+
+
+class KriegsmarineTestCase(unittest.TestCase):
+    """This is the Kriegsmarine example from Dirk Rijmenants' simulator manual.
+
+    It is credited to Stefan Krah and the M4 project:
+        http://www.bytereef.org/m4_project.html
+
+    """
+    def test_decrypt(self):
+
+        ring_settings = [ord(c) - ord('A') for c in 'AAAV']
+        stecker ='1/20 2/12 4/6 7/10 8/13 14/23 15/16 17/25 18/26 22/24'
+
+        machine = EnigmaMachine.from_key_sheet(
+                    rotors='Beta II IV I',
+                    ring_settings=ring_settings,
+                    reflector='B-Thin',
+                    plugboard_settings=stecker)
+
+        ciphertext = (
+        'FCLC QRKN NCZW VUSX PNYM INHZ XMQX SFWX WLKJ AHSH NMCO CCAK UQPM KCSM'
+        'HKSE INJU SBLK IOSX CKUB HMLL XCSJ USRR DVKO HULX WCCB GVLI YXEO AHXR'
+        'HKKF VDRE WEZL XOBA FGYU JQUK GRTV UKAM EURB VEKS UHHV OYHA BCJW MAKL'
+        'FKLM YFVN RIZR VVRT KOFD ANJM OLBG FFLE OPRG TFLV RHOW OPBE KVWM UQFM'
+        'PWPA RMFH AGKX IIBG FCLC QRKM VA'
+        ).replace(' ', '')
+
+        # remove the message indicators from the message (the first and last 2
+        # groups of the message -- it appears the last partial group 'VA' should
+        # be removed also)
+
+        ciphertext = ciphertext[8:]
+        ciphertext = ciphertext[:-10]
+
+        machine.set_display('VJNA')
+        plaintext = machine.process_text(ciphertext)
+
+        truth_data = (
+            'VONV ONJL OOKS JHFF TTTE'
+            'INSE INSD REIZ WOYY QNNS'
+            'NEUN INHA LTXX BEIA NGRI'
+            'FFUN TERW ASSE RGED RUEC'
+            'KTYW ABOS XLET ZTER GEGN'
+            'ERST ANDN ULAC HTDR EINU'
+            'LUHR MARQ UANT ONJO TANE'
+            'UNAC HTSE YHSD REIY ZWOZ'
+            'WONU LGRA DYAC HTSM YSTO'
+            'SSEN ACHX EKNS VIER MBFA'
+            'ELLT YNNN NNNO OOVI ERYS'
+            'ICHT EINS NULL'
+        ).replace(' ', '')
+
+        self.assertEqual(plaintext, truth_data)