#!/usr/bin/env python3
import sys
import base64
try:
    import click
except Exception as e:
    print(e)
    print("Install click")
    sys.exit(-1)


MAX_KEY_SIZE = 256


def printBrute(msg, after=None):
    for key in range(1, MAX_KEY_SIZE + 1):
        try:
            attempt = getTranslatedMessage(msg, key)
            print('[{:0>3}]'.format(key), after(attempt).decode())
        except:
            pass


def getTranslatedMessage(message, key):
    result = ""
    for char in message:
        if char.isalpha():
            lower_char = char.lower()
            new_ord = chr((ord(lower_char) - ord('a') + key) % 26 + ord('a'))
            if char == char.upper():
                new_ord = new_ord.upper()
            result += new_ord
        else:
            result += char
    return result


@click.command()
@click.option(
    "--key",
    "-k",
    nargs=1,
    type=click.IntRange(0, MAX_KEY_SIZE),
    default=13,
    help="encrypt message with `key`")
@click.option("--brute", "-b", is_flag=True, help="Bruteforce message")
@click.argument("message")
@click.option("--extra", "-e", nargs=1, type=click.Choice(["base64_after"]))
def cmd(message, brute, key=None, extra=None):
    if brute:
        if extra == "base64_after":
            printBrute(message, base64.b64decode)
    else:
        print(getTranslatedMessage(message, key))


if __name__ == '__main__':
    cmd()