Added some test cases and refactored keyword cipher out into a separate function
[cipher-tools.git] / cipher.py
index 0adfe904b6b388d069bf14231faa8e440fb67f7e..07b53926dbe1af98f1a34b1e76bb5e7286be0a1b 100644 (file)
--- a/cipher.py
+++ b/cipher.py
@@ -15,6 +15,10 @@ with open('count_1l.txt', 'r') as f:
         english_counts[letter] = int(count)
 normalised_english_counts = norms.normalise(english_counts)        
 
         english_counts[letter] = int(count)
 normalised_english_counts = norms.normalise(english_counts)        
 
+keywords = []
+with open('words.txt', 'r') as f:
+    keywords = [line.rstrip() for line in f]
+
 
 modular_division_table = [[0]*26 for x in range(26)]
 for a in range(26):
 
 modular_division_table = [[0]*26 for x in range(26)]
 for a in range(26):
@@ -68,6 +72,11 @@ def letter_frequencies(text):
         counts[c] += 1
     return counts
 
         counts[c] += 1
     return counts
 
+def deduplicate(text):
+    return list(collections.OrderedDict.fromkeys(text))
+
+
+
 def caesar_encipher_letter(letter, shift):
     """Encipher a letter, given a shift amount
 
 def caesar_encipher_letter(letter, shift):
     """Encipher a letter, given a shift amount
 
@@ -204,15 +213,64 @@ def affine_decipher(message, multiplier=1, adder=0, one_based=True):
     return ''.join(enciphered)
 
 
     return ''.join(enciphered)
 
 
+def keyword_cipher_alphabet_of(keyword, wrap_alphabet=False):
+    """Find the cipher alphabet given a keyword
+
+    >>> keyword_cipher_alphabet_of('harry')
+    'harybcdefgijklmnopqstuvwxz'
+    >>> keyword_cipher_alphabet_of('harry', True)
+    'haryzbcdefgijklmnopqstuvwx'
+    >>> keyword_cipher_alphabet_of('harry', False)
+    'harybcdefgijklmnopqstuvwxz'
+    """
+    cipher_alphabet = ''
+    if wrap_alphabet:
+        last_keyword_letter = deduplicate(sanitise(keyword))[-1]
+        last_keyword_position = string.ascii_lowercase.find(last_keyword_letter) + 1
+        cipher_alphabet = ''.join(deduplicate(sanitise(keyword) + string.ascii_lowercase[last_keyword_position:] + string.ascii_lowercase))
+    else:
+        cipher_alphabet = ''.join(deduplicate(sanitise(keyword) + string.ascii_lowercase))
+    return cipher_alphabet
+
+
+def keyword_encipher(message, keyword, wrap_alphabet=False):
+    """Enciphers a message with a keyword substitution cipher
+
+    >>> keyword_encipher('test message', 'harry')
+    'sbqs kbqqhdb'
+    >>> keyword_encipher('test message', 'harry', True)
+    'qzpq jzpphcz'
+    >>> keyword_encipher('test message', 'harry', False)
+    'sbqs kbqqhdb'
+    """
+    cipher_alphabet = keyword_cipher_alphabet_of(keyword, wrap_alphabet)
+    cipher_translation = ''.maketrans(string.ascii_lowercase, cipher_alphabet)
+    return message.lower().translate(cipher_translation)
+
+def keyword_decipher(message, keyword, wrap_alphabet=False):
+    """Deciphers a message with a keyword substitution cipher
+
+    >>> keyword_decipher('sbqs kbqqhdb', 'harry')
+    'test message'
+    >>> keyword_decipher('qzpq jzpphcz', 'harry', True)
+    'test message'
+    >>> keyword_decipher('sbqs kbqqhdb', 'harry', False)
+    'test message'
+    """
+    cipher_alphabet = keyword_cipher_alphabet_of(keyword, wrap_alphabet)
+    cipher_translation = ''.maketrans(cipher_alphabet, string.ascii_lowercase)
+    return message.lower().translate(cipher_translation)
+
+
 def caesar_break(message, metric=norms.euclidean_distance, target_frequencies=normalised_english_counts, message_frequency_scaling=norms.normalise):
     """Breaks a Caesar cipher using frequency analysis
     
 def caesar_break(message, metric=norms.euclidean_distance, target_frequencies=normalised_english_counts, message_frequency_scaling=norms.normalise):
     """Breaks a Caesar cipher using frequency analysis
     
-    >>> caesar_break('ibxcsyorsaqcheyklxivoexlevmrimwxsfiqevvmihrsasrxliwyrhecjsppsamrkwleppfmergefifvmhixscsymjcsyqeoixlm')
-    (4, 0.3186395289018361)
-    >>> caesar_break('wxwmaxdgheetgwuxztgptedbgznitgwwhpguxyhkxbmhvvtlbhgteeraxlmhiixweblmxgxwmhmaxybkbgztgwztsxwbgmxgmert')
-    (19, 0.4215290123583277)
-    >>> caesar_break('yltbbqnqnzvguvaxurorgenafsbezqvagbnornfgsbevpnaabjurersvaquvzyvxrnznazlybequrvfohgriraabjtbaruraprur')
-    (13, 0.31602920807545154)
+    >>> caesar_break('ibxcsyorsaqcheyklxivoexlevmrimwxsfiqevvmihrsasrxliwyrhecjsppsamrkwleppfmergefifvmhixscsymjcsyqeoixlm') # doctest: +ELLIPSIS
+    (4, 0.31863952890183...)
+    >>> caesar_break('wxwmaxdgheetgwuxztgptedbgznitgwwhpguxyhkxbmhvvtlbhgteeraxlmhiixweblmxgxwmhmaxybkbgztgwztsxwbgmxgmert') # doctest: +ELLIPSIS
+    (19, 0.42152901235832...)
+    >>> caesar_break('yltbbqnqnzvguvaxurorgenafsbezqvagbnornfgsbevpnaabjurersvaquvzyvxrnznazlybequrvfohgriraabjtbaruraprur') # doctest: +ELLIPSIS
+    (13, 0.316029208075451...)
     """
     sanitised_message = sanitise(message)
     best_shift = 0
     """
     sanitised_message = sanitise(message)
     best_shift = 0
@@ -221,7 +279,7 @@ def caesar_break(message, metric=norms.euclidean_distance, target_frequencies=no
         plaintext = caesar_decipher(sanitised_message, shift)
         frequencies = message_frequency_scaling(letter_frequencies(plaintext))
         fit = metric(target_frequencies, frequencies)
         plaintext = caesar_decipher(sanitised_message, shift)
         frequencies = message_frequency_scaling(letter_frequencies(plaintext))
         fit = metric(target_frequencies, frequencies)
-        logger.info('Caesar break attempt using key {0} gives fit of {1} and decrypt starting: {2}'.format(shift, fit, plaintext[:50]))
+        logger.debug('Caesar break attempt using key {0} gives fit of {1} and decrypt starting: {2}'.format(shift, fit, plaintext[:50]))
         if fit < best_fit:
             best_fit = fit
             best_shift = shift
         if fit < best_fit:
             best_fit = fit
             best_shift = shift
@@ -231,8 +289,8 @@ def caesar_break(message, metric=norms.euclidean_distance, target_frequencies=no
 def affine_break(message, metric=norms.euclidean_distance, target_frequencies=normalised_english_counts, message_frequency_scaling=norms.normalise):
     """Breaks an affine cipher using frequency analysis
     
 def affine_break(message, metric=norms.euclidean_distance, target_frequencies=normalised_english_counts, message_frequency_scaling=norms.normalise):
     """Breaks an affine cipher using frequency analysis
     
-    >>> affine_break('lmyfu bkuusd dyfaxw claol psfaom jfasd snsfg jfaoe ls omytd jlaxe mh jm bfmibj umis hfsul axubafkjamx. ls kffkxwsd jls ofgbjmwfkiu olfmxmtmwaokttg jlsx ls kffkxwsd jlsi zg tsxwjl. jlsx ls umfjsd jlsi zg hfsqysxog. ls dmmdtsd mx jls bats mh bkbsf. ls bfmctsd kfmyxd jls lyj, mztanamyu xmc jm clm cku tmmeaxw kj lai kxd clm ckuxj.')
-    ((15, 22, True), 0.2357036181865554)
+    >>> affine_break('lmyfu bkuusd dyfaxw claol psfaom jfasd snsfg jfaoe ls omytd jlaxe mh jm bfmibj umis hfsul axubafkjamx. ls kffkxwsd jls ofgbjmwfkiu olfmxmtmwaokttg jlsx ls kffkxwsd jlsi zg tsxwjl. jlsx ls umfjsd jlsi zg hfsqysxog. ls dmmdtsd mx jls bats mh bkbsf. ls bfmctsd kfmyxd jls lyj, mztanamyu xmc jm clm cku tmmeaxw kj lai kxd clm ckuxj.') # doctest: +ELLIPSIS
+    ((15, 22, True), 0.23570361818655...)
     """
     sanitised_message = sanitise(message)
     best_multiplier = 0
     """
     sanitised_message = sanitise(message)
     best_multiplier = 0
@@ -245,7 +303,7 @@ def affine_break(message, metric=norms.euclidean_distance, target_frequencies=no
                 plaintext = affine_decipher(sanitised_message, multiplier, adder, one_based)
                 frequencies = message_frequency_scaling(letter_frequencies(plaintext))
                 fit = metric(target_frequencies, frequencies)
                 plaintext = affine_decipher(sanitised_message, multiplier, adder, one_based)
                 frequencies = message_frequency_scaling(letter_frequencies(plaintext))
                 fit = metric(target_frequencies, frequencies)
-                logger.info('Affine break attempt using key {0}x+{1} ({2}) gives fit of {3} and decrypt starting: {4}'.format(multiplier, adder, one_based, fit, plaintext[:50]))
+                logger.debug('Affine break attempt using key {0}x+{1} ({2}) gives fit of {3} and decrypt starting: {4}'.format(multiplier, adder, one_based, fit, plaintext[:50]))
                 if fit < best_fit:
                     best_fit = fit
                     best_multiplier = multiplier
                 if fit < best_fit:
                     best_fit = fit
                     best_multiplier = multiplier
@@ -255,6 +313,29 @@ def affine_break(message, metric=norms.euclidean_distance, target_frequencies=no
     return (best_multiplier, best_adder, best_one_based), best_fit
 
 
     return (best_multiplier, best_adder, best_one_based), best_fit
 
 
+def keyword_break(message, wordlist=keywords, metric=norms.euclidean_distance, target_frequencies=normalised_english_counts, message_frequency_scaling=norms.normalise):
+    """Breaks a keyword substitution cipher using a dictionary and frequency analysis
+
+    >>> keyword_break(keyword_encipher('this is a test message for the keyword decipherment', 'elephant', True))
+    (('elephant', True), 0.41643991598441...) # doctest: +ELLIPSIS
+    """
+    best_keyword = ''
+    best_wrap_alphabet = True
+    best_fit = float("inf")
+    for wrap_alphabet in [True, False]:
+        for keyword in wordlist:
+            plaintext = keyword_decipher(message, keyword, wrap_alphabet)
+            frequencies = message_frequency_scaling(letter_frequencies(plaintext))
+            fit = metric(target_frequencies, frequencies)
+            logger.debug('Keyword break attempt using key {0} ({1}) gives fit of {2} and decrypt starting: {3}'.format(keyword, wrap_alphabet, fit, sanitise(plaintext)[:50]))
+            if fit < best_fit:
+                best_fit = fit
+                best_keyword = keyword
+                best_wrap_alphabet = wrap_alphabet
+    logger.info('Keyword break best fit with key {0} ({1}) gives fit of {2} and decrypt starting: {3}'.format(best_keyword, best_wrap_alphabet, best_fit, sanitise(keyword_decipher(message, best_keyword))[:50]))
+    return (best_keyword, best_wrap_alphabet), best_fit
+
+
 if __name__ == "__main__":
     import doctest
     doctest.testmod()
 if __name__ == "__main__":
     import doctest
     doctest.testmod()