vigniere.py

#!/usr/bin/python
# -*- coding: utf-8  -*-
import argparse
import math
import getpass
import sys

ALPHABET = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'

def test(orig, encr):
    o = ALPHABET.index(orig)
    e = ALPHABET.index(encr)
    c = e - o % 26
    if c < 0:
        c += 26
    return ALPHABET[c]

def grouped(s, block_width=5, blocks_per_row=4):
    groups = []
    for i in range(math.ceil(len(s) / block_width)):
        group = s[i * block_width:(i + 1) * block_width]
        groups.append(group)
    for i in range(math.ceil(len(groups) / blocks_per_row)):
        print(' '.join(groups[i * blocks_per_row:(i + 1) * blocks_per_row]))

def _code(content, code, mode):
    result = list(content)
    ext_code = [ALPHABET.index(c) for c in code]
    for i in range(len(content)):
        l = ALPHABET.find(content[i])
        if l >= 0:
            c = ext_code[len(ext_code) - len(code)]
            n, n_c = mode(l, c)
            n_c %= 26
            ext_code.append(n_c % 26)
            result[i] = ALPHABET[n % 26]
    return ''.join(result)

def decode(content, code):
    def _de(o, c):
        return [o - c]*2
    return _code(content, code, _de)

def encode(content, code):
    def _en(o, c):
        return o + c, o
    return _code(content, code, _en)

parser = argparse.ArgumentParser()
source = parser.add_mutually_exclusive_group(required=True)
source.add_argument('-i', '--inline')
source.add_argument('-f', '--file')
modes = parser.add_mutually_exclusive_group(required=True)
modes.add_argument('-e', '--encode', dest='mode', action='store_true')
modes.add_argument('-d', '--decode', dest='mode', action='store_false')
parser.add_argument('-v', '--verbose', action='store_true')
parser.add_argument('--save-commands', action='store_true')
parser.add_argument('-c', '--code')
args = parser.parse_args()

if args.file:
    with open(args.file, 'r') as f:
        args.inline = f.read()

if not args.save_commands:
    args.inline = [c for c in args.inline if c not in [' ', '\n', '\r']]

if not args.code:
    args.code = getpass.getpass()
    if not args.code:
        print("No code given, quitting.")
        sys.exit(1)

args.code = args.code.upper()
if not all(c in ALPHABET for c in args.code):
    print("Code may only contain letters.")
    sys.exit(1)

func = encode if args.mode else decode

if args.verbose:
    print('Input:')
    grouped(args.inline)
    print('Codeword:')
    print(args.code)
    print('Output:')
grouped(''.join(func(args.inline, args.code)))