3
\$\begingroup\$

I'm in the process of learning python and programmed this exercise in decrypting the Vigenere square cypher as practice.

Please comment on best practices, efficiency, or more pythonic ways to do things.

#!/usr/bin/env python

"""Functions for encrypting and decrypting text using
   the Vigenere square cipher. See:
   http://en.wikipedia.org/wiki/Vigen%C3%A8re_cipher

   IC stands for Index of Coincidence:
   http://en.wikipedia.org/wiki/Index_of_coincidence
"""

from __future__ import division
from collections import Counter
from math import fabs
from string import ascii_lowercase
from scipy.stats import pearsonr
from numpy import matrix
from os import system

#Define some constants:
LETTER_CNT = 26
ENGLISH_IC = 1.73

#Cornell English letter frequecy
#http://www.math.cornell.edu/~mec/2003-2004/cryptography/subs/frequencies.html
ENGLISH_LETTERS = 'etaoinsrhdlucmfywgpbvkxqjz'
ENGLISH_FREQ = [0.1202, 0.0910, 0.0812, 0.0768, 0.0731, 0.0695, 0.0628,
                0.0602, 0.0592, 0.0432, 0.0398, 0.0288, 0.0271, 0.0261,
                0.0230, 0.0211, 0.0209, 0.0203, 0.0182, 0.0149, 0.0111,
                0.0069, 0.0017, 0.0011, 0.0010, 0.0007]
ENGLISH_DICT = dict(zip(list(ENGLISH_LETTERS), ENGLISH_FREQ))
MAX_LEN = 10    #Maximum keyword length


def scrub_string(str):
    """Remove non-alphabetic characters and convert string to lower case. """
    return ''.join(ch for ch in str if ch.isalpha()).lower()


def string_to_numbers(str):
    """Convert str to a list of numbers giving the position of the letter
    in the alphabet (position of a = 0). str should contain only
    lowercase letters.
    """
    return [ord(ch) - ord('a') for ch in str]


def numbers_to_string(nums):
    """Convert a list of numbers to a string of letters
    (index of a = 0); the inverse function of string_to_numbers.
    """
    return ''.join(chr(n + ord('a')) for n in nums)


def shift_string_by_number(str, shift):
    """Shift the letters in str by the amount shift (either positive
    or negative) modulo 26.
    """
    return numbers_to_string((num + shift) % LETTER_CNT
                             for num in string_to_numbers(str))


def shift_string_by_letter(str, ch, direction):
    """Shift the letters in str by the value of ch, modulo 26.
    Right shift if direction = 1, left shift if direction = -1.
    """
    assert direction in {1, -1}
    return shift_string_by_number(str, (ord(ch) - ord('a') + 1) * direction)


def chunk_string(str):
    """Add a blank between each block of five characters in str."""
    return ' '.join(str[i:i+5] for i in xrange(0, len(str), 5))


def crypt(text, passphrase, which):
    """Encrypt or decrypt the text, depending on whether which = 1
    or which = -1.
    """
    text = scrub_string(text)
    passphrase = scrub_string(passphrase)
    letters = (shift_string_by_letter(ch, passphrase[i % len(passphrase)], which)
                   for i, ch in enumerate(text))
    return ''.join(letters)


def IC(text, ncol):
    """Divide the text into ncol columns and return the average index
    of coincidence across the columns.
    """
    text = scrub_string(text)
    A = str_to_matrix(scrub_string(text), ncol)
    cum = 0
    for col in A:
        N = len(col)
        cum += (sum(n*(n - 1) for n in Counter(col).values())
                / (N*(N - 1)/LETTER_CNT))
    return cum/ncol


def keyword_length(text):
    """Determine keyword length by finding the length that makes
    IC closest to the English plaintext value of 1.73.
    """
    text = scrub_string(text)
    a = [fabs(IC(text, ncol) - ENGLISH_IC) for ncol in xrange(1, MAX_LEN)]
    return a.index(min(a)) + 1


def correlation(letter_list):
    """Return the correlation of the frequencies of the letters
    in the list with the English letter frequency.
    """
    counts = Counter(letter_list)
    text_freq = [counts[ch]/len(letter_list) for ch in ascii_lowercase]
    english_freq = [ENGLISH_DICT[ch] for ch in ascii_lowercase]
    return pearsonr(text_freq, english_freq)[0]


def find_keyword_letter(letter_list):
    """Return a letter of the keyword, given every nth character
    of the ciphertext, where n = keyword length.
    """
    str = ''.join(letter_list)
    cors = [correlation(shift_string_by_number(str, -num))
            for num in xrange(1, LETTER_CNT + 1)]
    return ascii_lowercase[cors.index(max(cors))]


def find_keyword(ciphertext, keyword_length):
    """Return the keyword, given its length and the ciphertext."""
    A = str_to_matrix(scrub_string(ciphertext), keyword_length)
    return ''.join(
        [find_keyword_letter(A[j]) for j in xrange(keyword_length)])


def str_to_matrix(str, ncol):
    """Divide str into ncol lists as in the example below:

    >>> str_to_matrix('abcdefghijk', 4)
    [['a', 'e', 'i'], ['b', 'f', 'j'], ['c', 'g', 'k'], ['d', 'h']]
    """
    A = [list(str[i:i + ncol]) for i in xrange(0, len(str), ncol)]
    stub = A.pop()
    B = matrix(A).T.tolist()
    for i, ch in enumerate(stub):
        B[i] += ch
    return B


def test_functions():
    """Unit tests for functions in this module."""
    assert(shift_string_by_number('unladenswallow', 15) == 'jcapstchlpaadl')
    assert(shift_string_by_letter('unladenswallow', 'M', -1) == 'ngetwxglpteehp')
    assert(chunk_string('terpsichorean') == 'terps ichor ean')
    assert(crypt('Hello world!', "mypassword", 1) == 'udbmhplgdh')
    assert(crypt('udbmhplgdh', "mypassword", -1) == 'helloworld')
    assert(round(correlation('ganzunglabulich'), 6) == 0.118034)

    assert(scrub_string("I'm not Doctor bloody Bernofsky!!") ==
    'imnotdoctorbloodybernofsky')

    assert(string_to_numbers('lemoncurry') ==
    [11, 4, 12, 14, 13, 2, 20, 17, 17, 24])

    assert(numbers_to_string([11, 4, 12, 14, 13, 2, 20, 17, 17, 24]) ==
    'lemoncurry')

    assert(round(IC('''QPWKA LVRXC QZIKG RBPFA EOMFL JMSDZ VDHXC XJYEB IMTRQ
    WNMEA IZRVK CVKVL XNEIC FZPZC ZZHKM LVZVZ IZRRQ WDKEC
    HOSNY XXLSP MYKVQ XJTDC IOMEE XDQVS RXLRL KZHOV''', 5)
                 , 2) ==  1.82)

    assert(keyword_length('''QPWKA LVRXC QZIKG RBPFA EOMFL JMSDZ VDHXC XJYEB
    IMTRQ WNMEA IZRVK CVKVL XNEIC FZPZC ZZHKM LVZVZ IZRRQ WDKEC
    HOSNY XXLSP MYKVQ XJTDC IOMEE XDQVS RXLRL KZHOV''') == 5)

    assert(str_to_matrix('abcdefghijk', 4) ==
    [['a', 'e', 'i'], ['b', 'f', 'j'], ['c', 'g', 'k'], ['d', 'h']])


if __name__ == '__main__':
    print 'Calculating...'
    with open ("plaintext.txt", "r") as infile:
        plaintext = infile.read().replace('\n', ' ')
    passphrase = 'Moby Dick'
    ciphertext =  crypt(plaintext, passphrase, 1)
    kw_len = keyword_length(ciphertext)
    kw = find_keyword(ciphertext, kw_len)
    print 'Keyword length is {0}.'.format(kw_len)
    print 'The keyword is {0}.'.format(kw)
    system("""bash -c 'read -s -n 1 -p "Press any key print the decrypted text..."'""")
    print crypt(ciphertext, kw, -1)
\$\endgroup\$
3
\$\begingroup\$

Overall this is well documented, well written code. There are a number of things I may have written differently, but they are primarily manners of personal style. But I still found some things I want to call out that might be somewhat problematic, or at least worth examining:

  1. scrub_string might not do what you want. Notably, "\xe9".isalpha() is True with my default settings, but your code probably does not handle LATIN SMALL LETTER E WITH ACUTE.
  2. shift_string_by_letter seems like it might be better described as passing in the letter that a should become (or becomes that becomes a). As a related comment, make sure you are aware that assert lines are removed when python is run with -O or -OO so you cannot depend on them to catch run-time errors. Your use here is correct, as passing anything other than 1 or -1 is a programming error instead of a run-time error.
  3. crypt uses modulus to match what itertools could make simpler. With cycle and izip it can become this:

    letters = (shift_string_by_letter(ch, p, which)
                   for ch, p in izip(text, cycle(passphrase)))
    
  4. IC doesn't appear to follow the same naming conventions as your other functions. It's short, and upper-case, possibly an abbreviation.

  5. In __main__, the use of system and bash is unusual. I would replace this with raw_input("Press enter to print the decrypted text...") (or input(...) in python 3) to avoid the indirection that needlessly prevents this from working on systems without bash.
  6. In several places, constructs like ord(ch) - ord('a') or chr(n + ord('a')) help you convert between letters an indices. It might help either performance or readability to set up two dictionaries to do the conversion as a lookup, say to_index[ch] or from_index[n]. This would have a nice effect of catching unsupported characters more explicitly by raising a KeyError if one is encountered.
  7. As a general guideline, when you start to prioritize performance over readability, one of the lowest hanging fruits in python is function calls. While this can lead towards manually "inlining" other functions in order to save a few cycles, it can also lead towards other interesting approaches. For example, if the text for crypt is long enough, it's plausible that making up to 26 dictionaries mapping all 26 characters would allow for faster running code overall, despite the set-up time. The resulting code might look like this:

    shift_map = { p : { ch : shift_string_by_letter(ch, p, which)
                     for ch in ascii_lowercase } for p in passphrase }
    letters = (shift_map[p][ch] for ch, p in izip(text, cycle(passphrase)))
    

    But that strategy would fail utterly for short lengths of text due to the costs of setting up the dictionary. It's also possible that all of this would be lost in the noise of using generator expressions. This pontification highlights the need to profile the scenarios for which you actually want your code to perform well, and then to be creative about how you fix them.

\$\endgroup\$
  • 1
    \$\begingroup\$ This is great! I can't believe a site like this even exists, your efforts are much appreciated. \$\endgroup\$ – James King Dec 25 '13 at 17:48
2
\$\begingroup\$

You've written docstrings for all your functions, and test cases for many of them: this puts you ahead of 90% of the Python submissions here at Code Review! But there are still lots of small improvements you could make.

  1. Instead of:

    LETTER_CNT = 26
    

    I'd consider writing:

    LETTERS = ascii_lowercase
    LETTER_COUNT = len(LETTERS)
    

    to make it clear what is being counted. Also, there's no danger of running out of vowels, so I'd write COUNT out in full.

  2. By having separate tables ENGLISH_LETTERS and ENGLISH_FREQ and combining them using zip, you run the risk of getting these tables mismatched. It would be better to write:

    ENGLISH_DICT = dict(e=0.1202, t=0.0910, a=0.0812, ...)
    

    But all that you actually use this for is to build the following list in correlation:

    english_freq = [ENGLISH_DICT[ch] for ch in ascii_lowercase]
    

    So you could make this your global variable in the first place:

    # Frequency of letters a, b, c, ... in English.
    ENGLISH_FREQ = [0.0812, 0.0149, 0.0271, ...]
    

    and avoid building ENGLISH_DICT.

  3. This code is very nearly portable to Python 3. All that's remaining to do is to put parentheses around your print statements and change xrange to range.

  4. You name some of your parameters str, which is also the name of the built-in type str. This causes your variable to shadow the built-in, which would be inconvenient if you needed the built-in (as I do below).

  5. shift_string_by_number and shift_string_by_letter transform a string letter-wise. The built-in string method str.translate provides a simple and efficient way of doing this:

    def shift_string_by_number(s, shift):
        shift %= LETTER_COUNT
        tr = str.maketrans(LETTERS, LETTERS[shift:] + LETTERS[:shift])
        return s.translate(tr)
    

    For performance, it would make sense to pre-compute all the translation tables:

    SHIFT_TABLES = [str.maketrans(LETTERS, LETTERS[shift:] + LETTERS[:shift])
                    for shift in range(LETTER_COUNT)]
    
    def shift_string_by_number(s, shift):
        return s.translate(SHIFT_TABLES[shift % LETTER_COUNT])
    

    Notice that we only need 26 translation tables, even though shift might be negative, because the % operation always yields a result with the same sign as its second operand.

    After making this change, we no longer need string_to_numbers or numbers_to_string.

  6. The function scrub_string could also be rewritten to use str.translate. Here we can take advantage of the feature that "Characters mapped to None are deleted":

    from collections import defaultdict
    
    SCRUB_TABLE = defaultdict(lambda:None)
    SCRUB_TABLE.update(str.maketrans(LETTERS, LETTERS))
    SCRUB_TABLE.update(str.maketrans(ascii_uppercase, LETTERS))
    
    def scrub_string(s):
        """Remove non-alphabetic characters and convert string to lower case. """
        return s.translate(SCRUB_TABLE)
    
  7. You could benefit from rewriting your test cases to use the features in the unittest module. This would make it easier to run the tests (for example, by running python -munittest mymodule.py from the shell), but would also produce more informative results. A failed assertion just reports a bare AssertionError:

    >>> assert(shift_string_by_number('unladenswallow', 15) == 'jcapstchlpaadm')
    Traceback (most recent call last):
      File "<stdin>", line 1, in <module>
    AssertionError
    

    but if you used self.assertEqual in a unittest.TestCase, then you'd get output like this:

    $ python3.3 -munittest cr38055.py
    F
    ======================================================================
    FAIL: test_functions (cr38055.TestVigenere)
    Unit tests for functions in this module.
    ----------------------------------------------------------------------
    Traceback (most recent call last):
      File "./cr38055.py", line 149, in test_functions
        self.assertEqual(shift_string_by_number('unladenswallow', 15), 'jcapstchlpaadm')
    AssertionError: 'jcapstchlpaadl' != 'jcapstchlpaadm'
    - jcapstchlpaadl
    ?              ^
    + jcapstchlpaadm
    ?              ^
    
  8. In test_functions, you test crypt before you test scrub_string. But this means that an error in scrub_string will actually show up as an error in crypt, which will be harder for you to track down. It's best to organize your tests so that the functions are tested only after all their dependencies have been tested.

  9. The function crypt takes an argument which which needs to be −1 or +1 for encryption or decryption, but the docstring does not say which! (I think it's +1 for encryption and −1 for decryption.) Also, there's no check on the argument. It would be clearer to supply named constants:

    ENCRYPTION = +1
    DECRYPTION = -1
    

    and then check the argument:

    assert(which in (ENCRYPTION, DECRYPTION))
    

    or, better still, have two functions encrypt and decrypt.

  10. The function chunk_string appears not to be used.

  11. You call scrub_string at the start of keyword_length, and then that function then calls IC many times. Each call to IC therefore results in two useless calls to scrub_string. My suggestion would be to scrub strings just once (for example, when you read them from the source file) and not to scrub them inside worker functions like keyword_length, IC, crypt and so on.

  12. You use math.fabs to compute the absolute value of a floating-point number. But in fact the built-in abs would work just as well.

  13. In keyword_length you build a list a and then find the index of its minimum value:

    a = [fabs(IC(text, ncol) - ENGLISH_IC) for ncol in xrange(1, MAX_LEN)]
    return a.index(min(a)) + 1
    

    In Python it's rarely a good idea to bother with the index of an item in a sequence. Generally one should handle items directly and not via their index in a sequence.

    Note also that xrange(1, MAX_LEN) only goes up to MAX_LEN - 1 which seems wrong given that you documented MAX_LEN as "Maximum keyword length".

    You could fix the off-by-one error, avoid building the list, and avoid the call to index, by using the key argument to min, and writing:

    def keyword_length(text):
        def key(ncol):
            return abs(IC(text, ncol) - ENGLISH_IC)
        return min(range(1, MAX_LEN + 1), key=key)
    
  14. Similarly, in find_keyword_letter, you can use the key argument to max to avoid building a list and avoid a call to index:

    def find_keyword_letter(column):
        def key(letter):
            return correlation(shift_string_by_letter(column, letter, -1))
        return max(LETTERS, key=key)
    
  15. In IC, you could avoid the call to str_to_matrix by using Python's string slicing: the jth column from text is text[j::ncol]. Also, the multiplication by LETTER_CNT is common to all the summands so can be postponed to the end:

    def IC(text, ncol):
        total = 0
        for j in range(ncol):
            column = text[j::ncol]
            N = len(column)
            total += sum(n*(n-1) for n in Counter(column).values()) / (N*(N-1))
        return total * LETTER_COUNT / ncol
    
  16. In find_keyword you can avoid the call to str_to_matrix similarly:

    def find_keyword(ciphertext, ncol):
        return ''.join(find_keyword_letter(text[j::ncol]) for j in range(ncol))
    

    I've called the argument ncol here rather than keyword_length for consistency with the other functions, and so as not to shadow the function of that name.

    This eliminates the need for str_to_matrix or numpy.matrix.

\$\endgroup\$
  • \$\begingroup\$ Many excellent comments! I actually realized last night that I should not be calling a variable str and I've also ported to python 3. Thanks for responding! \$\endgroup\$ – James King Dec 29 '13 at 15:31

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.