#!/usr/bin/env python

from aesmix import MixSlice

import argparse
import logging
import os.path
import os

# enable logging
logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")


def fn_encrypt(args):
    key = args.key or os.urandom(16)
    iv = args.iv or os.urandom(16)
    output = args.output or (args.PLAINFILE + ".enc")
    public = args.public or (args.PLAINFILE + ".public")
    private = args.public or (args.PLAINFILE + ".private")

    with open(args.PLAINFILE, "rb") as fp:
        data = fp.read()

    logging.info("[*] Encrypting file %s ..." % args.PLAINFILE)
    manager = MixSlice.encrypt(data, key, iv)
    manager.save_to_files(output, public, private)
    logging.info("Output fragdir: %s" % output)
    logging.info("Public key file:  %s" % public)
    logging.info("Private key file: %s" % private)


def fn_update(args):
    public = args.public or (os.path.splitext(args.FRAGDIR)[0] + ".public")
    private = args.public or (os.path.splitext(args.FRAGDIR)[0] + ".private")
    logging.info("[*] Performing policy update on %s ..." % args.FRAGDIR)
    manager = MixSlice.load_from_file(args.FRAGDIR, private)
    manager.step_encrypt()
    manager.save_to_files(args.FRAGDIR, public, private)
    logging.info("Done")


def fn_decrypt(args):
    base = os.path.splitext(args.FRAGDIR)[0]
    keyfile = (args.public if args.public else
               args.private if args.private else
               base + ".public" if os.path.isfile(base + ".public") else
               base + ".private")
    assert os.path.isfile(keyfile), "keyfile not found."

    logging.info("[*] Decrypting fragdir %s using key %s ..." %
                 (args.FRAGDIR, keyfile))
    output = args.output or (args.FRAGDIR + ".dec")
    manager = MixSlice.load_from_file(args.FRAGDIR, keyfile)
    plaindata = manager.decrypt()
    with open(output, "wb") as fp:
        fp.write(plaindata)
    logging.info("Decrypted file: %s" % output)


def main():
    parser = argparse.ArgumentParser(
        description="Mix&Slice tool for encrypting/decrypting files",
        epilog="More info at: https://doi.org/10.1145/2976749.2978377")
    subparsers = parser.add_subparsers(help="sub-command help")

    keyparser = argparse.ArgumentParser(add_help=False)
    keyparser.add_argument("--public", default=None,
                           help="public key file (default: $INPUT.public)")
    keyparser.add_argument("--private", default=None,
                           help="private key file (default: $INPUT.private)")

    # parser for the encrypt command
    encrypt = subparsers.add_parser("encrypt", parents=[keyparser],
                                    help="encrypt a file")
    encrypt.add_argument("PLAINFILE", help="path to input file")
    encrypt.add_argument("--key", default=None, help="encryption key (16B)")
    encrypt.add_argument("--iv", default=None, help="encryption iv (16B)")
    encrypt.add_argument("--output", default=None,
                         help="output file (default: $INPUT.enc)")
    encrypt.set_defaults(func=fn_encrypt)

    # parser for the update command
    update = subparsers.add_parser("update", parents=[keyparser],
                                   help="perform policy update")
    update.add_argument("FRAGDIR", help="path to the fragments directory")
    update.set_defaults(func=fn_update)

    # parser for the decrypt command
    decrypt = subparsers.add_parser("decrypt", parents=[keyparser],
                                    help="decrypt a file")
    decrypt.add_argument("FRAGDIR", help="path to the fragments directory")
    decrypt.add_argument("--output", default=None,
                         help="output file (default: $FRAGDIR.dec)")
    decrypt.set_defaults(func=fn_decrypt)

    # parse args
    args = parser.parse_args()

    # normalize args
    try:
        args.key = args.key.encode("ascii")
    except Exception:
        pass

    try:
        args.iv = args.iv.encode("ascii")
    except Exception:
        pass

    # invoke function
    args.func(args)


if __name__ == "__main__":
    main()

# vim: set syntax=python:
