#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# This software is hereby released into public domain. Use it wisely.
#
# Originally written by Mirek Kratochvil (2017)
# Python3 port by Bernhard Esslinger (Feb 2018)

import argparse
import random
import sys

letters = "_ABCDEFGHIJKLMNOPQRSTUVWXYZ.,:?!'()0123456789+-*/"
tiles = list(zip(letters, map(lambda x: (x // 7, x % 7), range(7 * 7))))
padding_size_min = 5
padding_size_max = 15


def check_key(key):
    if len(key) != len(letters):
        raise ValueError("Wrong key size")
    cnts = {}
    for c in letters:
        cnts[c] = 0
    for c in key:
        if not c in cnts:
            raise ValueError(f"Letter {c} not in LS47!")
        cnts[c] += 1
        if cnts[c] > 1:
            raise ValueError(f"Letter {c} duplicated in key!")


def find_ix(letter):
    m = [l for l in tiles if l[0] == letter]
    if len(m) != 1:
        raise ValueError(f"Letter {letter} not in LS47!")
    for l, pos in m:
        return pos


def find_pos(key, letter):
    p = key.find(letter)
    if p >= 0 and p < 7 * 7:
        return (p // 7, p % 7)
    raise ValueError(f"Letter {letter} not in key!")


def add_pos(a, b):
    return ((a[0] + b[0]) % 7, (a[1] + b[1]) % 7)


def sub_pos(a, b):
    return ((a[0] - b[0]) % 7, (a[1] - b[1]) % 7)


def find_at_pos(key, coord):
    return key[coord[1] + coord[0] * 7]


def rotate_right(key, row, n):
    mid = key[7 * row : 7 * (row + 1)]
    n = (7 - n % 7) % 7
    return key[: 7 * row] + mid[n:] + mid[:n] + key[7 * (row + 1) :]


def rotate_down(key, col, n):
    lines = [key[i * 7 : (i + 1) * 7] for i in range(7)]
    lefts = [l[:col] for l in lines]
    mids = [l[col] for l in lines]
    rights = [l[col + 1 :] for l in lines]
    n = (7 - n % 7) % 7
    mids = mids[n:] + mids[:n]
    return "".join(lefts[i] + mids[i] + rights[i] for i in range(7))


def rotate_marker_right(m, row, n):
    if m[0] != row:
        return (m[0], m[1])
    else:
        return (m[0], (m[1] + n) % 7)


def rotate_marker_down(m, col, n):
    if m[1] != col:
        return (m[0], m[1])
    else:
        return ((m[0] + n) % 7, m[1])


def derive_key(password):
    k = ''.join(reversed(letters))
    i = 0
    for c in password + ' ':
        ix = k.find(c,i)
        if ix<0:
            continue
        k = k[:i] + k[ix] + k[i:ix] + k[ix+1:]
        i += 1
    return ''.join(reversed(k))


def encrypt(key, plaintext):
    check_key(key)
    mp = (0, 0)
    ciphertext = ""
    for p in plaintext:
        pp = find_pos(key, p)
        mix = find_ix(find_at_pos(key, mp))
        cp = add_pos(pp, mix)
        c = find_at_pos(key, cp)
        ciphertext += c

        key = rotate_right(key, pp[0], 1)
        cp = find_pos(key, c)
        key = rotate_down(key, cp[1], 1)
        mp = add_pos(mp, find_ix(c))
    return ciphertext


def decrypt(key, ciphertext):
    check_key(key)
    mp = (0, 0)
    plaintext = ""
    for c in ciphertext:
        cp = find_pos(key, c)
        mix = find_ix(find_at_pos(key, mp))
        pp = sub_pos(cp, mix)
        p = find_at_pos(key, pp)
        plaintext += p

        key = rotate_right(key, pp[0], 1)
        cp = find_pos(key, c)
        key = rotate_down(key, cp[1], 1)
        mp = add_pos(mp, find_ix(c))
    return plaintext


def encrypt_pad(key, plaintext, signature):
    check_key(key)
    padding = "".join(
        map(
            lambda x: letters[random.randint(1, len(letters) - 1)],
            range(random.randint(padding_size_min, padding_size_max)),
        )
    )

    return encrypt(key, padding + "_" + plaintext + "_" + signature)


def decrypt_pad(key, ciphertext, signature):
    check_key(key)
    padded = decrypt(key, ciphertext)
    pad_end = padded.find("_")
    if pad_end < padding_size_min or pad_end > padding_size_max:
        raise ValueError("mangled padding")
    padded = padded[pad_end + 1 :]
    if len(signature) == 0:
        return padded
    if padded[-len(signature) - 1 :] != "_" + signature:
        raise ValueError("bad signature")
    return padded[: -len(signature) - 1]


if __name__ == "__main__":

    ap = argparse.ArgumentParser(
        prog=sys.argv[0],
        description="Computer implementation of LS47.",
        epilog="This is flimsy! Encrypt physically!",
    )

    ap.add_argument(
        "-d", "--decrypt", dest="decrypt", action="store_true", help="run decryption"
    )
    ap.add_argument(
        "-p",
        "--pad-and-sign",
        dest="sign",
        help="On encryption, add a random padding and a signature at the end. On decryption, expect a random padding and the exact signature. Specify empty string to attach no (or accept any) signature.",
    )
    ap.add_argument(
        "-s",
        "--password",
        dest="password",
        help="Password string (to be expanded by the simple swapping algorithm)",
    )
    ap.add_argument(
        "-K",
        "--full-key",
        dest="fullkey",
        help="Shared key (the initial state of the board, line by line, always with `_' instead of space; used only if password is not specified)",
    )
    ap.add_argument(
        "-I",
        "--exact-input",
        dest="inx",
        help="Do not convert input spaces to `_', ignore them instead.",
        action="store_true"
    )
    ap.add_argument(
        "-O",
        "--exact-output",
        dest="outx",
        help="Do not convert output `_'s to spaces.",
        action="store_true"
    )

    args = ap.parse_args()

    if args.password:
        key = derive_key(args.password if args.inx else args.password.replace(' ', '_'))
    elif args.fullkey:
        key = args.fullkey
    else:
        raise ValueError("Specify at least one of password and full key (options -s and -K)")

    input = sys.stdin.read()
    if not args.inx:
        input = input.replace(' ', '_')
    letter_set = set(letters)
    input = ''.join(filter(lambda x:x in letter_set, input.upper()))

    if args.decrypt:
        if args.sign is None:
            output = decrypt(key, input)
        else:
            output = decrypt_pad(key, input, args.sign)
    else:  # encrypt
        if args.sign is None:
            output = encrypt(key, input)
        else:
            output = encrypt_pad(key, input, args.sign)

    if not args.outx:
        output = output.replace('_', ' ')

    print(output)