vigenere-test.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

# Created by Michal Wiglasz <michalwiglasz.cz>
# Feel free to use, but mention my name somewhere.


from difflib import SequenceMatcher
import subprocess
import glob
import os


def _wrap_with(code):
    def inner(text, bold=False):
        c = code
        if bold:
            c = "1;%s" % c
        return "\033[%sm%s\033[0m" % (c, text)
    return inner


red = _wrap_with('31')
green = _wrap_with('32')
yellow = _wrap_with('33')
blue = _wrap_with('34')
magenta = _wrap_with('35')
cyan = _wrap_with('36')
white = _wrap_with('37')


def colorize(key, correct_key):
    s = SequenceMatcher(a=key, b=correct_key)
    diff = list(key) or ['']
    for tag, i1, i2, j1, j2 in s.get_opcodes():
        if tag == 'replace':
            diff[i1:i2] = map(yellow, diff[i1:i2])

        elif tag == 'insert':
            if (i1 == 0):
                diff[i1] = green(correct_key[j1:j2]) + diff[i1]
            else:
                diff[i1 - 1] += green(correct_key[j1:j2])

        elif tag == 'delete':
            diff[i1:i2] = map(red, diff[i1:i2])

    return ''.join(diff)


def main():
    passed = 0
    failed = 0
    total = 0
    wrong_length = 0

    for name in sorted(glob.iglob('tests/*.cip')):
        with open(name, 'rt') as cipherfile:
            correct_key = os.path.basename(name)[:-4]
            output = subprocess.check_output(['./kry'], stdin=cipherfile).strip()
            friedman, kasiski, length, key = output.split(';')

            if (key == correct_key):
                print green("[X]"), " %s == %s (%s)" % (correct_key, key, output)
                passed += 1

            else:
                print red("[!]"), " %s != %s (%s)" % (correct_key, colorize(key, correct_key), output)
                failed += 1
                if len(key) != len(correct_key):
                    wrong_length += 1

            total += 1

    print ""

    if failed == 0:
        print green("COOL. Passed all %d tests" % (total,))

    elif passed == 0:
        print red("SCHEIßE. Failed all %d tests" % (total,))

    else:
        print yellow("Passed %d/%d tests" % (passed, total))
        print yellow("%d wrong lengths, %d wrong letters" % (wrong_length, failed - wrong_length))


if __name__ == '__main__':
    main()