Data commit

Task/Bifid-cipher/Python/bifid-cipher.py

@@ -0,0 +1,106 @@
+"""Bifid cipher. Requires Python >=3.7."""
+import math
+import pprint
+import string
+
+from itertools import chain
+from itertools import zip_longest
+
+from typing import Dict
+from typing import Iterable
+from typing import Iterator
+from typing import Tuple
+from typing import TypeVar
+
+
+T = TypeVar("T")
+
+
+def group(it: Iterable[T], n: int) -> Iterator[Tuple[T, ...]]:
+    """Return the input iterable split in to `n` equal chunks, padded with `None`."""
+    return zip_longest(*[iter(it)] * n)
+
+
+Square = Tuple[Tuple[str, ...], ...]
+
+
+def polybius_square(alphabet: str) -> Square:
+    """Return the given alphabet as a tuple of tuples, representing a Polybius square."""
+    return tuple(group(alphabet, math.ceil(math.sqrt(len(alphabet)))))
+
+
+def polybius_map(square: Square) -> Dict[str, Tuple[int, int]]:
+    """Return a reverse lookup for the given Polybius square."""
+    return {
+        square[i][j]: (i + 1, j + 1)
+        for i in range(len(square))
+        for j in range(len(square))
+    }
+
+
+def encrypt(message: str, square: Square) -> str:
+    """Encrypt a plaintext message using a bifid cipher with the given Polybius square."""
+    _map = polybius_map(square)
+    return "".join(
+        square[x - 1][y - 1]
+        for x, y in group(
+            chain.from_iterable(zip(*(_map[c] for c in message if c in _map))),
+            2,
+        )
+    )
+
+
+def decrypt(message: str, square: Square) -> str:
+    """Decrypt a ciphertext message using a bifid cipher with the given Polybius square."""
+    _map = polybius_map(square)
+    return "".join(
+        square[x - 1][y - 1]
+        for x, y in zip(
+            *group(
+                chain.from_iterable((_map[c] for c in message if c in _map)),
+                len(message),
+            )
+        )
+    )
+
+
+def normalize(message: str) -> str:
+    """Normalize a message for the typical Polybius square."""
+    return message.upper().replace("J", "I")
+
+
+TYPICAL_POLYBIUS_SQUARE = polybius_square(
+    alphabet="".join(c for c in string.ascii_uppercase if c != "J")
+)
+
+
+EXAMPLE_POLYBIUS_SQUARE = polybius_square(
+    alphabet="BGWKZQPNDSIOAXEFCLUMTHYVR",
+)
+
+
+def main() -> None:
+    test_cases = [
+        ("ATTACKATDAWN", TYPICAL_POLYBIUS_SQUARE),  # 1
+        ("FLEEATONCE", EXAMPLE_POLYBIUS_SQUARE),  # 2
+        ("FLEEATONCE", TYPICAL_POLYBIUS_SQUARE),  # 3
+        (
+            normalize("The invasion will start on the first of January"),
+            polybius_square(alphabet="PLAYFIREXMBCDGHKNOQSTUVWZ"),
+        ),
+        (
+            "The invasion will start on the first of January".upper(),
+            polybius_square(alphabet=string.ascii_uppercase + string.digits),
+        ),
+    ]
+
+    for message, square in test_cases:
+        pprint.pprint(square)
+        print("Message  :", message)
+        print("Encrypted:", encrypt(message, square))
+        print("Decrypted:", decrypt(encrypt(message, square), square))
+        print("")
+
+
+if __name__ == "__main__":
+    main()