caesar.py

import re
import sys
import random
import argparse


if sys.argv[1] == 'generate':
    generate = True
else:
    generate = False
    parser = argparse.ArgumentParser()
    parser.add_argument('-d', '--decode',
            help='decode the message', action='store_true')
    parser.add_argument('-e', '--encode',
            help='encode the message (default)', action='store_true')
    parser.add_argument('message', help='message to encode/decode')
    parser.add_argument('key', help='encryption key')
    args = parser.parse_args()


def normalize_key(key):
    if key.isdigit() and not len(key) % 2:
        new_key = []
        for key_num in re.findall('..', key):
            key_num = int(key_num)
            while key_num > 26:
                key_num -= 26
            new_key.append(key_num)
        return new_key
    else:
        raise parser.error('key argument must be an even number of integers')


def rotate(message, key, mode):
    if mode == 'd':
        key = [-x for x in key]
    elif mode != 'e':
        raise ValueError('mode argument must be "d" or "e"')

    for rot_num in key:
        translated = ''
        for char in message:
            if char.isalpha():
                num = ord(char)
                num += rot_num
                if char.isupper():
                    if num > ord('Z'):
                        num -= 26
                    elif num < ord('A'):
                        num += 26
                elif char.islower():
                    if num > ord('z'):
                        num -= 26
                    elif num < ord('a'):
                        num += 26
                translated += chr(num)
            else:
                translated += char
        message = translated
        x = 12

    return translated


if generate:
    key = [str(random.randint(1, 25)).zfill(2) for x in range(5)]
    print(''.join(key))
else:
    key = normalize_key(args.key)
    mode = 'd' if args.decode else 'e'
    print(rotate(args.message, key, mode))