diff --git a/imp/crypto/cyphers.py b/imp/crypto/cyphers.py index b4b1297..b6896eb 100644 --- a/imp/crypto/cyphers.py +++ b/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], - ignore_noncharset: bool = True) -> str: - cyphertext = '' + charsets: list[str] | list[list] = CHARSET_SUB_DEFAULT, + ignore_noncharset: bool = True, + 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: - raise ValueError(ERRMSG_ROT_CHARSET_DUPL.format(c)) - break + if active_charset is not None: + raise ValueError(ERRMSG_CHARSET_DUPL.format('ROT', c)) + 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 caesar(plaintext: str) -> str: return ROT(plaintext, 1) +def ROT13(plaintext: str, charsets: list[str] = CHARSET_SUB_DEFAULT) -> str: + 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 +