Fixed spelling of KeywordWrapAlphabet
[cipher-training.git] / cipher.py
index ba62d411a902e40f6bc1b259d2733111f891407d..f5c5d33f412108c045f47effc1e93c221e90c118 100644 (file)
--- a/cipher.py
+++ b/cipher.py
@@ -27,7 +27,7 @@ def every_nth(text, n, fillvalue=''):
     >>> every_nth(string.ascii_lowercase, 5, fillvalue='!')
     ['afkpuz', 'bglqv!', 'chmrw!', 'dinsx!', 'ejoty!']
     """
-    split_text = [text[i:i+n] for i in range(0, len(text), n)]
+    split_text = chunks(text, n, fillvalue)
     return [''.join(l) for l in zip_longest(*split_text, fillvalue=fillvalue)]
 
 def combine_every_nth(split_text):
@@ -212,7 +212,7 @@ def affine_decipher_letter(letter, multiplier=1, adder=0, one_based=True):
         if one_based: cipher_number += 1
         plaintext_number = ( 
             modular_division_table[multiplier]
-                                  [(cipher_number - adder) % 26] )
+                                  [(cipher_number - adder) % 26])
         if one_based: plaintext_number -= 1
         return chr(plaintext_number % 26 + alphabet_start) 
     else:
@@ -241,31 +241,31 @@ def affine_decipher(message, multiplier=1, adder=0, one_based=True):
     return ''.join(enciphered)
 
 
-class Keyword_wrap_alphabet(Enum):
-    from_a = 0
-    from_last = 1
-    from_largest = 2
+class KeywordWrapAlphabet(Enum):
+    from_a = 1
+    from_last = 2
+    from_largest = 3
 
 
-def keyword_cipher_alphabet_of(keyword, wrap_alphabet=Keyword_wrap_alphabet.from_a):
+def keyword_cipher_alphabet_of(keyword, wrap_alphabet=KeywordWrapAlphabet.from_a):
     """Find the cipher alphabet given a keyword.
     wrap_alphabet controls how the rest of the alphabet is added
     after the keyword.
 
     >>> keyword_cipher_alphabet_of('bayes')
     'bayescdfghijklmnopqrtuvwxz'
-    >>> keyword_cipher_alphabet_of('bayes', Keyword_wrap_alphabet.from_a)
+    >>> keyword_cipher_alphabet_of('bayes', KeywordWrapAlphabet.from_a)
     'bayescdfghijklmnopqrtuvwxz'
-    >>> keyword_cipher_alphabet_of('bayes', Keyword_wrap_alphabet.from_last)
+    >>> keyword_cipher_alphabet_of('bayes', KeywordWrapAlphabet.from_last)
     'bayestuvwxzcdfghijklmnopqr'
-    >>> keyword_cipher_alphabet_of('bayes', Keyword_wrap_alphabet.from_largest)
+    >>> keyword_cipher_alphabet_of('bayes', KeywordWrapAlphabet.from_largest)
     'bayeszcdfghijklmnopqrtuvwx'
     """
-    if wrap_alphabet == Keyword_wrap_alphabet.from_a:
+    if wrap_alphabet == KeywordWrapAlphabet.from_a:
         cipher_alphabet = ''.join(deduplicate(sanitise(keyword) + 
                                               string.ascii_lowercase))
     else:
-        if wrap_alphabet == Keyword_wrap_alphabet.from_last:
+        if wrap_alphabet == KeywordWrapAlphabet.from_last:
             last_keyword_letter = deduplicate(sanitise(keyword))[-1]
         else:
             last_keyword_letter = sorted(sanitise(keyword))[-1]
@@ -278,7 +278,7 @@ def keyword_cipher_alphabet_of(keyword, wrap_alphabet=Keyword_wrap_alphabet.from
     return cipher_alphabet
 
 
-def keyword_encipher(message, keyword, wrap_alphabet=Keyword_wrap_alphabet.from_a):
+def keyword_encipher(message, keyword, wrap_alphabet=KeywordWrapAlphabet.from_a):
     """Enciphers a message with a keyword substitution cipher.
     wrap_alphabet controls how the rest of the alphabet is added
     after the keyword.
@@ -288,18 +288,18 @@ def keyword_encipher(message, keyword, wrap_alphabet=Keyword_wrap_alphabet.from_
 
     >>> keyword_encipher('test message', 'bayes')
     'rsqr ksqqbds'
-    >>> keyword_encipher('test message', 'bayes', Keyword_wrap_alphabet.from_a)
+    >>> keyword_encipher('test message', 'bayes', KeywordWrapAlphabet.from_a)
     'rsqr ksqqbds'
-    >>> keyword_encipher('test message', 'bayes', Keyword_wrap_alphabet.from_last)
+    >>> keyword_encipher('test message', 'bayes', KeywordWrapAlphabet.from_last)
     'lskl dskkbus'
-    >>> keyword_encipher('test message', 'bayes', Keyword_wrap_alphabet.from_largest)
+    >>> keyword_encipher('test message', 'bayes', KeywordWrapAlphabet.from_largest)
     'qspq jsppbcs'
     """
     cipher_alphabet = keyword_cipher_alphabet_of(keyword, wrap_alphabet)
     cipher_translation = ''.maketrans(string.ascii_lowercase, cipher_alphabet)
     return unaccent(message).lower().translate(cipher_translation)
 
-def keyword_decipher(message, keyword, wrap_alphabet=Keyword_wrap_alphabet.from_a):
+def keyword_decipher(message, keyword, wrap_alphabet=KeywordWrapAlphabet.from_a):
     """Deciphers a message with a keyword substitution cipher.
     wrap_alphabet controls how the rest of the alphabet is added
     after the keyword.
@@ -309,11 +309,11 @@ def keyword_decipher(message, keyword, wrap_alphabet=Keyword_wrap_alphabet.from_
     
     >>> keyword_decipher('rsqr ksqqbds', 'bayes')
     'test message'
-    >>> keyword_decipher('rsqr ksqqbds', 'bayes', Keyword_wrap_alphabet.from_a)
+    >>> keyword_decipher('rsqr ksqqbds', 'bayes', KeywordWrapAlphabet.from_a)
     'test message'
-    >>> keyword_decipher('lskl dskkbus', 'bayes', Keyword_wrap_alphabet.from_last)
+    >>> keyword_decipher('lskl dskkbus', 'bayes', KeywordWrapAlphabet.from_last)
     'test message'
-    >>> keyword_decipher('qspq jsppbcs', 'bayes', Keyword_wrap_alphabet.from_largest)
+    >>> keyword_decipher('qspq jsppbcs', 'bayes', KeywordWrapAlphabet.from_largest)
     'test message'
     """
     cipher_alphabet = keyword_cipher_alphabet_of(keyword, wrap_alphabet)
@@ -471,15 +471,18 @@ def scytale_encipher(message, rows, fillvalue=' '):
     >>> scytale_encipher('thequickbrownfox', 4)
     'tubnhirfecooqkwx'
     >>> scytale_encipher('thequickbrownfox', 5)
-    'tubnhirfecooqkwx'
+    'tubn hirf ecoo qkwx '
     >>> scytale_encipher('thequickbrownfox', 6)
     'tqcrnxhukof eibwo '
     >>> scytale_encipher('thequickbrownfox', 7)
-    'tqcrnxhukof eibwo '
+    'tqcrnx hukof  eibwo  '
     """
-    transpositions = [i for i in range(math.ceil(len(message) / rows))]
+    # transpositions = [i for i in range(math.ceil(len(message) / rows))]
+    # return column_transposition_encipher(message, transpositions, 
+    #     fillvalue=fillvalue, fillcolumnwise=False, emptycolumnwise=True)
+    transpositions = [i for i in range(rows)]
     return column_transposition_encipher(message, transpositions, 
-        fillcolumnwise=False, emptycolumnwise=True)
+        fillvalue=fillvalue, fillcolumnwise=True, emptycolumnwise=False)
 
 def scytale_decipher(message, rows):
     """Deciphers using the scytale transposition cipher.
@@ -489,18 +492,180 @@ def scytale_decipher(message, rows):
     'thequickbrownfox  '
     >>> scytale_decipher('tubnhirfecooqkwx', 4)
     'thequickbrownfox'
-    >>> scytale_decipher('tubnhirfecooqkwx', 5)
-    'thequickbrownfox'
+    >>> scytale_decipher('tubn hirf ecoo qkwx ', 5)
+    'thequickbrownfox    '
     >>> scytale_decipher('tqcrnxhukof eibwo ', 6)
     'thequickbrownfox  '
-    >>> scytale_decipher('tqcrnxhukof eibwo ', 7)
-    'thequickbrownfox  '
+    >>> scytale_decipher('tqcrnx hukof  eibwo  ', 7)
+    'thequickbrownfox     '
     """
-    transpositions = [i for i in range(math.ceil(len(message) / rows))]
+    # transpositions = [i for i in range(math.ceil(len(message) / rows))]
+    # return column_transposition_decipher(message, transpositions, 
+    #     fillcolumnwise=False, emptycolumnwise=True)
+    transpositions = [i for i in range(rows)]
     return column_transposition_decipher(message, transpositions, 
-        fillcolumnwise=False, emptycolumnwise=True)
+        fillcolumnwise=True, emptycolumnwise=False)
+
+
+class PocketEnigma(object):
+    """A pocket enigma machine
+    The wheel is internally represented as a 26-element list self.wheel_map, 
+    where wheel_map[i] == j shows that the position i places on from the arrow 
+    maps to the position j places on.
+    """
+    def __init__(self, wheel=1, position='a'):
+        """initialise the pocket enigma, including which wheel to use and the
+        starting position of the wheel.
+
+        The wheel is either 1 or 2 (the predefined wheels) or a list of letter
+        pairs.
+
+        The position is the letter pointed to by the arrow on the wheel.
+
+        >>> pe.wheel_map
+        [25, 4, 23, 10, 1, 7, 9, 5, 12, 6, 3, 17, 8, 14, 13, 21, 19, 11, 20, 16, 18, 15, 24, 2, 22, 0]
+        >>> pe.position
+        0
+        """
+        self.wheel1 = [('a', 'z'), ('b', 'e'), ('c', 'x'), ('d', 'k'), 
+            ('f', 'h'), ('g', 'j'), ('i', 'm'), ('l', 'r'), ('n', 'o'), 
+            ('p', 'v'), ('q', 't'), ('s', 'u'), ('w', 'y')]
+        self.wheel2 = [('a', 'c'), ('b', 'd'), ('e', 'w'), ('f', 'i'), 
+            ('g', 'p'), ('h', 'm'), ('j', 'k'), ('l', 'n'), ('o', 'q'), 
+            ('r', 'z'), ('s', 'u'), ('t', 'v'), ('x', 'y')]
+        if wheel == 1:
+            self.make_wheel_map(self.wheel1)
+        elif wheel == 2:
+            self.make_wheel_map(self.wheel2)
+        else:
+            self.validate_wheel_spec(wheel)
+            self.make_wheel_map(wheel)
+        if position in string.ascii_lowercase:
+            self.position = ord(position) - ord('a')
+        else:
+            self.position = position
+
+    def make_wheel_map(self, wheel_spec):
+        """Expands a wheel specification from a list of letter-letter pairs
+        into a full wheel_map.
+
+        >>> pe.make_wheel_map(pe.wheel2)
+        [2, 3, 0, 1, 22, 8, 15, 12, 5, 10, 9, 13, 7, 11, 16, 6, 14, 25, 20, 21, 18, 19, 4, 24, 23, 17]
+        """
+        self.validate_wheel_spec(wheel_spec)
+        self.wheel_map = [0] * 26
+        for p in wheel_spec:
+            self.wheel_map[ord(p[0]) - ord('a')] = ord(p[1]) - ord('a')
+            self.wheel_map[ord(p[1]) - ord('a')] = ord(p[0]) - ord('a')
+        return self.wheel_map
+
+    def validate_wheel_spec(self, wheel_spec):
+        """Validates that a wheel specificaiton will turn into a valid wheel
+        map.
+
+        >>> pe.validate_wheel_spec([])
+        Traceback (most recent call last):
+            ...
+        ValueError: Wheel specification has 0 pairs, requires 13
+        >>> pe.validate_wheel_spec([('a', 'b', 'c')]*13)
+        Traceback (most recent call last):
+            ...
+        ValueError: Not all mappings in wheel specificationhave two elements
+        >>> pe.validate_wheel_spec([('a', 'b')]*13)
+        Traceback (most recent call last):
+            ...
+        ValueError: Wheel specification does not contain 26 letters
+        """
+        if len(wheel_spec) != 13:
+            raise ValueError("Wheel specification has {} pairs, requires 13".
+                format(len(wheel_spec)))
+        for p in wheel_spec:
+            if len(p) != 2:
+                raise ValueError("Not all mappings in wheel specification"
+                    "have two elements")
+        if len(set([p[0] for p in wheel_spec] + 
+                    [p[1] for p in wheel_spec])) != 26:
+            raise ValueError("Wheel specification does not contain 26 letters")
+
+    def encipher_letter(self, letter):
+        """Enciphers a single letter, by advancing the wheel before looking up
+        the letter on the wheel.
+
+        >>> pe.set_position('f')
+        5
+        >>> pe.encipher_letter('k')
+        'h'
+        """
+        self.advance()
+        return self.lookup(letter)
+    decipher_letter = encipher_letter
+
+    def lookup(self, letter):
+        """Look up what a letter enciphers to, without turning the wheel.
+
+        >>> pe.set_position('f')
+        5
+        >>> ''.join([pe.lookup(l) for l in string.ascii_lowercase])
+        'udhbfejcpgmokrliwntsayqzvx'
+        >>> pe.lookup('A')
+        ''
+        """
+        if letter in string.ascii_lowercase:
+            return chr(
+                (self.wheel_map[(ord(letter) - ord('a') - self.position) % 26] + 
+                    self.position) % 26 + 
+                ord('a'))
+        else:
+            return ''
+
+    def advance(self):
+        """Advances the wheel one position.
+
+        >>> pe.set_position('f')
+        5
+        >>> pe.advance()
+        6
+        """
+        self.position = (self.position + 1) % 26
+        return self.position
+
+    def encipher(self, message, starting_position=None):
+        """Enciphers a whole message.
+
+        >>> pe.set_position('f')
+        5
+        >>> pe.encipher('helloworld')
+        'kjsglcjoqc'
+        >>> pe.set_position('f')
+        5
+        >>> pe.encipher('kjsglcjoqc')
+        'helloworld'
+        >>> pe.encipher('helloworld', starting_position = 'x')
+        'egrekthnnf'
+        """
+        if starting_position:
+            self.set_position(starting_position)
+        transformed = ''
+        for l in message:
+            transformed += self.encipher_letter(l)
+        return transformed
+    decipher = encipher
+
+    def set_position(self, position):
+        """Sets the position of the wheel, by specifying the letter the arrow
+        points to.
+
+        >>> pe.set_position('a')
+        0
+        >>> pe.set_position('m')
+        12
+        >>> pe.set_position('z')
+        25
+        """
+        self.position = ord(position) - ord('a')
+        return self.position
 
 
 if __name__ == "__main__":
     import doctest
-    doctest.testmod()
+    doctest.testmod(extraglobs={'pe': PocketEnigma(1, 'a')})