Wrote 6b solution
[cipher-tools.git] / cipher.py
index 1c71b5c205567807332accf8d3fbada74b1c0c28..86125b38e805654f0957ec5e4fd27d7e285c5177 100644 (file)
--- a/cipher.py
+++ b/cipher.py
@@ -13,6 +13,7 @@ import pprint
 ## Utility functions
 cat = ''.join
 wcat = ' '.join
+lcat = '\n'.join
 
 def pos(letter): 
     if letter in string.ascii_lowercase:
@@ -381,7 +382,7 @@ def vigenere_encipher(message, keyword):
     >>> vigenere_encipher('hello', 'abc')
     'hfnlp'
     """
-    shifts = [ord(l) - ord('a') for l in sanitise(keyword)]
+    shifts = [pos(l) for l in sanitise(keyword)]
     pairs = zip(message, cycle(shifts))
     return cat([caesar_encipher_letter(l, k) for l, k in pairs])
 
@@ -391,12 +392,129 @@ def vigenere_decipher(message, keyword):
     >>> vigenere_decipher('hfnlp', 'abc')
     'hello'
     """
-    shifts = [ord(l) - ord('a') for l in sanitise(keyword)]
+    shifts = [pos(l) for l in sanitise(keyword)]
     pairs = zip(message, cycle(shifts))
     return cat([caesar_decipher_letter(l, k) for l, k in pairs])
 
-beaufort_encipher=vigenere_decipher
-beaufort_decipher=vigenere_encipher
+
+def beaufort_encipher(message, keyword):
+    """Beaufort encipher
+
+    >>> beaufort_encipher('inhisjournaldatedtheidesofoctober', 'arcanaimperii')
+    'sevsvrusyrrxfayyxuteemazudmpjmmwr'
+    """
+    shifts = [pos(l) for l in sanitise(keyword)]
+    pairs = zip(message, cycle(shifts))
+    return cat([unpos(k - pos(l)) for l, k in pairs])
+
+beaufort_decipher = beaufort_encipher    
+
+beaufort_variant_encipher=vigenere_decipher
+beaufort_variant_decipher=vigenere_encipher
+
+
+def polybius_grid(keyword, column_order, row_order, letters_to_merge=None,
+                  wrap_alphabet=KeywordWrapAlphabet.from_a):
+    """Grid for a Polybius cipher, using a keyword to rearrange the
+    alphabet.
+
+
+    >>> polybius_grid('a', 'abcde', 'abcde')['x'] == ('e', 'c')
+    True
+    >>> polybius_grid('elephant', 'abcde', 'abcde')['e'] == ('a', 'a')
+    True
+    >>> polybius_grid('elephant', 'abcde', 'abcde')['b'] == ('b', 'c')
+    True
+    """
+    alphabet = keyword_cipher_alphabet_of(keyword, wrap_alphabet=wrap_alphabet)
+    if letters_to_merge is None: 
+        letters_to_merge = {'j': 'i'}
+    grid = {l: k 
+            for k, l in zip([(c, r) for c in column_order for r in row_order],
+                [l for l in alphabet if l not in letters_to_merge])}
+    for l in letters_to_merge:
+        grid[l] = grid[letters_to_merge[l]]
+    return grid
+
+def polybius_reverse_grid(keyword, column_order, row_order, letters_to_merge=None,
+                  wrap_alphabet=KeywordWrapAlphabet.from_a):
+    """Grid for decrypting using a Polybius cipher, using a keyword to 
+    rearrange the alphabet.
+
+    >>> polybius_reverse_grid('a', 'abcde', 'abcde')['e', 'c'] == 'x'
+    True
+    >>> polybius_reverse_grid('elephant', 'abcde', 'abcde')['a', 'a'] == 'e'
+    True
+    >>> polybius_reverse_grid('elephant', 'abcde', 'abcde')['b', 'c'] == 'b'
+    True
+    """
+    alphabet = keyword_cipher_alphabet_of(keyword, wrap_alphabet=wrap_alphabet)
+    if letters_to_merge is None: 
+        letters_to_merge = {'j': 'i'}
+    grid = {k: l 
+            for k, l in zip([(c, r) for c in column_order for r in row_order],
+                [l for l in alphabet if l not in letters_to_merge])}
+    return grid  
+
+
+def polybius_flatten(pair, column_first):
+    """Convert a series of pairs into a single list of characters"""
+    if column_first:
+        return str(pair[1]) + str(pair[0])
+    else:
+        return str(pair[0]) + str(pair[1])
+
+def polybius_encipher(message, keyword, column_order, row_order, 
+                      column_first=False,
+                      letters_to_merge=None, wrap_alphabet=KeywordWrapAlphabet.from_a): 
+    """Encipher a message with Polybius cipher, using a keyword to rearrange
+    the alphabet
+
+
+    >>> polybius_encipher('this is a test message for the ' \
+          'polybius decipherment', 'elephant', \
+          [1, 2, 3, 4, 5], [1, 2, 3, 4, 5], \
+          wrap_alphabet=KeywordWrapAlphabet.from_last)
+    '2214445544551522115522511155551543114252542214111352123234442355411135441314115451112122'
+    >>> polybius_encipher('this is a test message for the ' \
+          'polybius decipherment', 'elephant', 'abcde', 'abcde', \
+          column_first=False)
+    'bbadccddccddaebbaaddbbceaaddddaecbaacadadcbbadaaacdaabedbcccdeddbeaabdccacadaadcceaababb'
+    >>> polybius_encipher('this is a test message for the ' \
+          'polybius decipherment', 'elephant', 'abcde', 'abcde', \
+          column_first=True)
+    'bbdaccddccddeabbaaddbbecaaddddeabcaaacadcdbbdaaacaadbadecbccedddebaadbcccadaaacdecaaabbb'
+    """
+    grid = polybius_grid(keyword, column_order, row_order, letters_to_merge, wrap_alphabet)
+    return cat(polybius_flatten(grid[l], column_first)
+               for l in message
+               if l in grid)
+
+
+def polybius_decipher(message, keyword, column_order, row_order, 
+                      column_first=False,
+                      letters_to_merge=None, wrap_alphabet=KeywordWrapAlphabet.from_a):    
+    """Decipher a message with a Polybius cipher, using a keyword to rearrange
+    the alphabet
+
+    >>> polybius_decipher('bbdaccddccddeabbaaddbbecaaddddeabcaaacadcdbbdaaaca'\
+    'adbadecbccedddebaadbcccadaaacdecaaabbb', 'elephant', 'abcde', 'abcde', \
+    column_first=False)
+    'toisisvtestxessvbephktoefhnugiysweqifoekxelt'
+
+    >>> polybius_decipher('bbdaccddccddeabbaaddbbecaaddddeabcaaacadcdbbdaaaca'\
+    'adbadecbccedddebaadbcccadaaacdecaaabbb', 'elephant', 'abcde', 'abcde', \
+    column_first=True)
+    'thisisatestmessageforthepolybiusdecipherment'
+    """
+    grid = polybius_reverse_grid(keyword, column_order, row_order, letters_to_merge, wrap_alphabet)
+    column_index_type = type(column_order[0])
+    row_index_type = type(row_order[0])
+    if column_first:
+        pairs = [(column_index_type(p[1]), row_index_type(p[0])) for p in chunks(message, 2)]
+    else:
+        pairs = [(row_index_type(p[0]), column_index_type(p[1])) for p in chunks(message, 2)]
+    return cat(grid[p] for p in pairs if p in grid)
 
 
 def transpositions_of(keyword):
@@ -786,12 +904,12 @@ def hill_encipher(matrix, message_letters, fillvalue='a'):
         padding = fillvalue[0] * (n - len(sanitised_message) % n)
     else:
         padding = ''
-    message = [ord(c) - ord('a') for c in sanitised_message + padding]
+    message = [pos(c) for c in sanitised_message + padding]
     message_chunks = [message[i:i+n] for i in range(0, len(message), n)]
     # message_chunks = chunks(message, len(matrix), fillvalue=None)
     enciphered_chunks = [((matrix * np.matrix(c).T).T).tolist()[0] 
             for c in message_chunks]
-    return cat([chr(int(round(l)) % 26 + ord('a')) 
+    return cat([unpos(round(l))
             for l in sum(enciphered_chunks, [])])
 
 def hill_decipher(matrix, message, fillvalue='a'):
@@ -1053,7 +1171,7 @@ class PocketEnigma(object):
             self.validate_wheel_spec(wheel)
             self.make_wheel_map(wheel)
         if position in string.ascii_lowercase:
-            self.position = ord(position) - ord('a')
+            self.position = pos(position)
         else:
             self.position = position
 
@@ -1067,8 +1185,8 @@ class PocketEnigma(object):
         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')
+            self.wheel_map[pos(p[0])] = pos(p[1])
+            self.wheel_map[pos(p[1])] = pos(p[0])
         return self.wheel_map
 
     def validate_wheel_spec(self, wheel_spec):
@@ -1123,10 +1241,9 @@ class PocketEnigma(object):
         ''
         """
         if letter in string.ascii_lowercase:
-            return chr(
-                (self.wheel_map[(ord(letter) - ord('a') - self.position) % 26] + 
-                    self.position) % 26 + 
-                ord('a'))
+            return unpos(
+                (self.wheel_map[(pos(letter) - self.position) % 26] + 
+                    self.position))
         else:
             return ''
 
@@ -1174,7 +1291,7 @@ class PocketEnigma(object):
         >>> pe.set_position('z')
         25
         """
-        self.position = ord(position) - ord('a')
+        self.position = pos(position)
         return self.position