diff --git a/ciphers/two_square_cipher.py b/ciphers/two_square_cipher.py new file mode 100644 index 000000000000..2070c973f232 --- /dev/null +++ b/ciphers/two_square_cipher.py @@ -0,0 +1,108 @@ +def generate_square(keyword: str) -> list: + """ + Generate a 5x5 Playfair square (I/J combined) from a keyword. + """ + keyword = keyword.upper().replace("J", "I") + seen = set() + square = [] + for char in keyword: + if char.isalpha() and char not in seen: + seen.add(char) + square.append(char) + for char in "ABCDEFGHIKLMNOPQRSTUVWXYZ": + if char not in seen: + seen.add(char) + square.append(char) + return [square[i : i + 5] for i in range(0, 25, 5)] + + +def find_position(square, char): + """ + Return (row, col) of char in given 5x5 square. + """ + if char == "J": + char = "I" + for r in range(5): + for c in range(5): + if square[r][c] == char: + return r, c + return None + + +def two_square_encrypt(plaintext: str, key1: str, key2: str) -> str: + """ + Encrypt plaintext using the Two-Square cipher. + + >>> two_square_encrypt("HELLO", "KEYWORD", "EXAMPLE") + 'BFCCLN' + """ + plaintext = plaintext.upper().replace("J", "I") + plaintext = "".join([c for c in plaintext if c.isalpha()]) + + # split into digraphs + pairs = [] + i = 0 + while i < len(plaintext): + a = plaintext[i] + b = plaintext[i + 1] if i + 1 < len(plaintext) else "X" + if a == b: + b = "X" + i += 1 + else: + i += 2 + pairs.append((a, b)) + + square1 = generate_square(key1) + square2 = generate_square(key2) + + ciphertext = "" + for a, b in pairs: + r1, c1 = find_position(square1, a) + r2, c2 = find_position(square2, b) + ciphertext += square2[r1][c2] + ciphertext += square1[r2][c1] + return ciphertext + + +def two_square_decrypt(ciphertext: str, key1: str, key2: str) -> str: + """ + Decrypt ciphertext using the Two-Square cipher. + + >>> two_square_decrypt("BFCCLN", "KEYWORD", "EXAMPLE") + 'HELLOX' + """ + ciphertext = ciphertext.upper().replace("J", "I") + ciphertext = "".join([c for c in ciphertext if c.isalpha()]) + + if len(ciphertext) % 2 != 0: + raise ValueError("Ciphertext must have even length") + + square1 = generate_square(key1) + square2 = generate_square(key2) + + plaintext = "" + for i in range(0, len(ciphertext), 2): + a, b = ciphertext[i], ciphertext[i + 1] + r1, c1 = find_position(square2, a) + r2, c2 = find_position(square1, b) + plaintext += square1[r1][c2] + plaintext += square2[r2][c1] + return plaintext + + +if __name__ == "__main__": + from doctest import testmod + + testmod() + + # Example usage + plaintext = "HELLO" + key1 = "KEYWORD" + key2 = "EXAMPLE" + encrypted = two_square_encrypt(plaintext, key1, key2) + decrypted = two_square_decrypt(encrypted, key1, key2) + + print("\n\n") + print("Plaintext:", plaintext) + print("Encrypted:", encrypted) + print("Decrypted:", decrypted)