Using ngram length for scytale break

diff --git a/cipher.py b/cipher.py
index ba4a73f7f39ab8d712b9ba9e52d49e269e1999ed..0e22a2a85bbc0afa5c6dc1b5caaf6c80459fcb2e 100644 (file)
--- a/cipher.py
+++ b/cipher.py
@@ -655,11 +655,12 @@ def scytale_break(message,
     """
     best_key = 0
     best_fit = float("inf")
+    ngram_length = len(next(iter(target_counts.keys())))
     for key in range(1, 20):
         if len(message) % key == 0:
             plaintext = scytale_decipher(message, key)
             counts = message_frequency_scaling(frequencies(
-                         ngrams(sanitise(plaintext), 2)))
+                         ngrams(sanitise(plaintext), ngram_length)))
             fit = metric(target_counts, counts)
             logger.debug('Scytale break attempt using key {0} gives fit of '
                          '{1} and decrypt starting: {2}'.format(key,