Playfair cipher

ciphers/playfair_cipher.py

@@ -0,0 +1,102 @@
+import string
+import itertools
+
+def chunker(seq, size):
+    it = iter(seq)
+    while True:
+       chunk = tuple(itertools.islice(it, size))
+       if not chunk:
+           return
+       yield chunk
+
+
+
+def prepare_input(dirty):
+    """
+    Prepare the plaintext by uppcasing it 
+    and seperating repeated letters with X's
+    """
+
+    dirty = ''.join([c.upper() for c in dirty if c in string.ascii_letters])
+    clean = ""
+
+    if len(dirty) < 2:
+        return dirty
+
+    for i in range(len(dirty)-1):
+        clean += dirty[i]
+
+        if dirty[i] == dirty[i+1]:
+            clean += 'X'
+
+    clean += dirty[-1]
+
+    if len(clean) & 1:
+        clean += 'X'
+
+    return clean
+
+def generate_table(key):
+
+    # I and J are used interchangably to allow
+    # us to use a 5x5 table (25 letters)
+    alphabet = "ABCDEFGHIKLMNOPQRSTUVWXYZ"
+    # we're using a list instead of a '2d' array because it makes the math 
+    # for setting up the table and doing the actual encoding/decoding simpler
+    table = []
+
+    # copy key chars into the table if they are in `alphabet` ignoring duplicates
+    for char in key.upper():
+        if char not in table and char in alphabet:
+            table.append(char)
+
+    # fill the rest of the table in with the remaining alphabet chars
+    for char in alphabet:
+        if char not in table:
+            table.append(char)
+
+    return table
+
+def encode(plaintext, key):
+    table = generate_table(key)
+    plaintext = prepare_input(plaintext)
+    ciphertext = ""
+
+    # https://en.wikipedia.org/wiki/Playfair_cipher#Description
+    for char1, char2 in chunker(plaintext, 2):
+        row1, col1 = divmod(table.index(char1), 5)
+        row2, col2 = divmod(table.index(char2), 5)
+
+        if row1 == row2:
+            ciphertext += table[row1*5+(col1+1)%5]
+            ciphertext += table[row2*5+(col2+1)%5]
+        elif col1 == col2:
+            ciphertext += table[((row1+1)%5)*5+col1]
+            ciphertext += table[((row2+1)%5)*5+col2]
+        else: # rectangle
+            ciphertext += table[row1*5+col2]
+            ciphertext += table[row2*5+col1]
+
+    return ciphertext
+
+
+def decode(ciphertext, key):
+    table = generate_table(key)
+    plaintext = ""
+
+    # https://en.wikipedia.org/wiki/Playfair_cipher#Description
+    for char1, char2 in chunk(ciphertext, 2):
+        row1, col1 = divmod(table.index(char1), 5)
+        row2, col2 = divmod(table.index(char2), 5)
+
+        if row1 == row2:
+            plaintext += table[row1*5+(col1-1)%5]
+            plaintext += table[row2*5+(col2-1)%5]
+        elif col1 == col2:
+            plaintext += table[((row1-1)%5)*5+col1]
+            plaintext += table[((row2-1)%5)*5+col2]
+        else: # rectangle
+            plaintext += table[row1*5+col2]
+            plaintext += table[row2*5+col1]
+
+    return plaintext