codebook.substitution

  1"""
  2This module defines a number of classical substitution ciphers.
  3
  4It features different types of substitution ciphers:
  5
  6  - Monoalphabetic: `caesar`, `generic`, `keyphrase`
  7  - Polyalphabetic: `vigenere`
  8  - Polygraphic: `playfair`
  9"""
 10import itertools
 11from string import ascii_lowercase, ascii_uppercase
 12
 13from codebook.utils import (
 14    codegroup,
 15    validate_cipher_alphabet,
 16    validate_key,
 17    validate_plaintext,
 18)
 19
 20
 21def _shifted_alphabet(n: int) -> str:
 22    return ascii_uppercase[n:] + ascii_uppercase[:n]
 23
 24
 25def _keyed_alphabet(key: str, *, from_start: bool) -> str:
 26    seen = set(key)
 27
 28    if from_start:
 29        remaining = [c for c in ascii_uppercase if c not in seen]
 30    else:
 31        start = ord(key[-1]) - 65
 32        remaining = [c for c in _shifted_alphabet(start + 1) if c not in seen]
 33
 34    return key + "".join(remaining)
 35
 36
 37@codegroup
 38def caesar(plaintext: str, *, shift: int) -> str:
 39    """Caesar cipher (page 10)
 40
 41    - `plaintext` is the message to be encrypted
 42    - `shift` is the number of places to rotate the plain alphabet (right shift)
 43    """
 44    alphabet = _shifted_alphabet(shift % 26)
 45    return generic(plaintext, cipher_alphabet=alphabet)
 46
 47
 48@codegroup
 49def generic(plaintext: str, *, cipher_alphabet: str) -> str:
 50    """Monoalphabetic substitution cipher (page 12)
 51
 52    - `plaintext` is the message to be encrypted
 53    - `cipher_alphabet` is the cipher alphabet that represents the substitution
 54    """
 55    plaintext = validate_plaintext(plaintext)
 56    cipher_alphabet = validate_cipher_alphabet(cipher_alphabet)
 57
 58    mapping = str.maketrans(ascii_lowercase, cipher_alphabet)
 59
 60    return plaintext.translate(mapping)
 61
 62
 63@codegroup
 64def keyphrase(plaintext: str, *, key: str) -> str:
 65    """Monoalphabetic substitution cipher using a keyword/keyphrase (page 13)
 66
 67    - `plaintext` is the message to be encrypted
 68    - `key` is the keyword or keyphrase used to generate the cipher alphabet
 69
 70    As per the example in the book, the remainder of the alphabet starts
 71    where the keyphrase ends.
 72    For example, `JULIUS CAESAR` generates the `JULISCAERTVWXYZBDFGHKMNOPQ` alphabet
 73    and not `JULISCAERBDFGHKMNOPQTVWXYZ`, which is a common alternative.
 74    """
 75    key = validate_key(key)
 76
 77    alphabet = _keyed_alphabet(key, from_start=False)
 78    return generic(plaintext, cipher_alphabet=alphabet)
 79
 80
 81@codegroup
 82def vigenere(plaintext: str, *, key: str) -> str:
 83    """Vigenère cipher (page 48)
 84
 85    - `plaintext` is the message to be encrypted
 86    - `key` defines the series of interwoven Caesar ciphers to be used
 87    """
 88    plaintext = validate_plaintext(plaintext)
 89    key = validate_key(key)
 90
 91    cycled_cipher_alphabet = itertools.cycle(
 92        _shifted_alphabet(ord(c) - 65) for c in key
 93    )
 94
 95    seq = [next(cycled_cipher_alphabet)[ord(c) - 97] for c in plaintext]
 96    return "".join(seq)
 97
 98
 99@codegroup
100def playfair(plaintext: str, *, key: str) -> str:
101    """Playfair cipher (Appendix E, page 372)
102
103    - `plaintext` is the message to be encrypted
104    - `key` defines the 5x5 table used for encryption
105
106    As per the example in the book, `J` and `I` share a space on the table
107    and `x` is used as padding for *same-letter* digrams and for the *last space*,
108    if required.
109    """
110    plaintext = validate_plaintext(plaintext)
111    key = validate_key(key)
112
113    # build matrix
114    cipher_alphabet = list(_keyed_alphabet(key, from_start=True))
115    cipher_alphabet.remove("J")
116
117    from_char, to_char = {}, {}
118    for i, c in enumerate(cipher_alphabet):
119        y, x = divmod(i, 5)
120        from_char[c] = (x, y)
121        to_char[x, y] = c
122
123    from_char["J"] = from_char["I"]
124
125    # collect digraphs
126    digraphs = []
127    i = 0
128    while i + 1 < len(plaintext):
129        a, b = plaintext[i], plaintext[i + 1]
130        if a != b:
131            digraphs.append(a + b)
132            i += 2
133        else:
134            digraphs.append(a + "x")
135            i += 1
136
137    if i < len(plaintext):
138        digraphs.append(plaintext[i] + "x")
139
140    # apply transfrom
141    seq = []
142    for digraph in digraphs:
143        x, y = from_char[digraph[0].upper()]
144        v, w = from_char[digraph[1].upper()]
145
146        if y == w:
147            a = to_char[(x + 1) % 5, y]
148            b = to_char[(v + 1) % 5, y]
149        elif x == v:
150            a = to_char[x, (y + 1) % 5]
151            b = to_char[x, (w + 1) % 5]
152        else:
153            a = to_char[v, y]
154            b = to_char[x, w]
155
156        seq.append(a + b)
157
158    return "".join(seq)