implemented multiple simple mono/polyalphabetic substitution ciphers

imp/crypto/cyphers.py

@@ -1,32 +1,137 @@
+'''
+Implementation of various substitution cyphers, some of which are 
+already in the standard library but lack features/options I require.
+
+Terminology:
+1. "Simple": simple substitution cyphers operates on single letters
+2. "Polygraphic": polygraphic substitution cyphers operate on groups
+                  of letters larger than length 1.
+3. "Monoalphabetic cypher": a monoalphabetic cypher applies the same fixed
+                            substitutions to the entire message (ie ROT13)
+4. "Polyalphabetic cypher": a polyalphabetic cypher applies different 
+                            substitutions to the entire message (ie vigenere)
+'''
+
 from imp.constants import ALPHA_LOWER, ALPHA_UPPER
 
-ERRMSG_ROT_CHARSET_DUPL = 'Bad charsets for ROT: found duplicate character \'{0}\''
+'''
-ERRMSG_ROT_CHAR_UNKNOWN = 'No charset for ROT has character \'{0}\' (exception thrown since `ignore_noncharset=False`)'
+Constant Declarations 
+'''
+# Error Message Format Strings
+ERRMSG_CHARSET_DUPL = 'Bad charsets for {0}: found duplicate character \'{1}\''
+ERRMSG_NOT_IN_CHARSET = 'No charset for {0} has character \'{1}\''
+# Charsets
+CHARSET_SUB_DEFAULT = [ALPHA_LOWER, ALPHA_UPPER]
 
 '''
-Substitution Cyphers  
+==========================================
+Simple Monoalphabetic Substitution Cyphers
+==========================================
 '''
-def ROT(plaintext: str, 
+def ROT(plaintext: str | list, 
         n: int, 
-        charsets: list[str] = [ALPHA_LOWER, ALPHA_UPPER], 
+        charsets: list[str] | list[list] = CHARSET_SUB_DEFAULT, 
-        ignore_noncharset: bool = True) -> str:
+        ignore_noncharset: bool = True,
-    cyphertext = ''
+        as_string: bool = True) -> str:
+    cyphertext = []
     for c in plaintext:
-        index = None
+        i = None
+        active_charset = None
         for charset in charsets:
             try:
                 i = charset.index(c)
                 # if we reached here then no ValueError occured
-                if index is not None:
+                if active_charset is not None:
-                    raise ValueError(ERRMSG_ROT_CHARSET_DUPL.format(c))
+                    raise ValueError(ERRMSG_CHARSET_DUPL.format('ROT', c))
-                break
+                active_charset = charset
             except ValueError: pass
-        if index is not None:
+            
-            cyphertext += charset[(i + n) % len(charset)]
+        if i is not None:
+            cyphertext.append(active_charset[(i + n) % len(active_charset)])
         elif not ignore_noncharset:
-            raise ValueError(ERRMSG_ROT_CHAR_UNKNOWN.format(c))
+            raise ValueError(ERRMSG_NOT_IN_CHARSET.format('ROT', c))
+        else:
+            # if ignore_noncharset then just append c
+            cyphertext.append(c)
+    if as_string:
+        return ''.join(cyphertext)
     return cyphertext
 
 # Common ROT aliases
-def ROT13(plaintext: str) -> str: return ROT(plaintext, 13)
+def ROT13(plaintext: str, charsets: list[str] = CHARSET_SUB_DEFAULT) -> str: 
-def caesar(plaintext: str) -> str: return ROT(plaintext, 1)
+    return ROT(plaintext, 13, charsets=charsets)
+def caesar(plaintext: str, charsets: list[str] = CHARSET_SUB_DEFAULT) -> str: 
+    return ROT(plaintext, 1, charsets=charsets)
+
+# The Atbash Cypher maps a character to its reverse (ie a -> z, b -> y, etc)
+def atbash(plaintext: str | list, 
+           charsets: list[str] | list[list] = CHARSET_SUB_DEFAULT,
+           ignore_noncharset: bool = True,
+           as_string: bool = True) -> str:
+    cyphertext = []
+    for c in plaintext:
+        i = None
+        active_charset = None
+        for charset in charsets:
+            try:
+                i = charset.index(c)
+                # if we reached here then no ValueError occured
+                if active_charset is not None:
+                    raise ValueError(ERRMSG_CHARSET_DUPL.format('Atbash', c))
+                active_charset = charset
+            except ValueError: pass
+            
+        if i is not None:
+            L = len(active_charset)
+            cyphertext.append(active_charset[(L - i - 1) % L])
+        elif not ignore_noncharset:
+            raise ValueError(ERRMSG_NOT_IN_CHARSET.format('Atbash', c))
+        else:
+            # if ignore_noncharset then just append c
+            cyphertext.append(c)
+    if as_string:
+        return ''.join(cyphertext)
+    return cyphertext
+
+'''
+=========================================
+Simple Polyalphabetic Substituion Cyphers
+========================================= 
+'''
+def vigenere(plaintext: str | list, 
+             key: str | list, 
+             charsets: list[str] | list[list] = CHARSET_SUB_DEFAULT,
+             key_charset: str | list = ALPHA_UPPER,
+             decrypt: bool = False,
+             ignore_noncharset: bool = True,
+             as_string: bool = True) -> str:
+    cyphertext = []
+    # map the key characters to their rotations
+    rots = [key_charset.index(k) for k in key]
+    pos = 0
+    for c in plaintext:
+        i = None
+        active_charset = None
+        for charset in charsets:
+            try:
+                i = charset.index(c)
+                # if we reached here then no ValueError occured
+                if active_charset is not None:
+                    raise ValueError(ERRMSG_CHARSET_DUPL.format('Vigenere', c))
+                active_charset = charset
+            except ValueError: pass
+            
+        if i is not None:
+            rot = rots[pos % len(rots)]
+            if decrypt: rot = -rot
+            cyphertext.append(active_charset[(i + rot) % len(active_charset)])
+            pos += 1
+        elif not ignore_noncharset:
+            raise ValueError(ERRMSG_NOT_IN_CHARSET.format('Vigenere', c))
+        else:
+            # if ignore_noncharset then just append c
+            cyphertext.append(c)
+    if as_string:
+        return ''.join(cyphertext)
+    return cyphertext
+