# Copyright (c) 2015 Egor Tensin <Egor.Tensin@gmail.com>
# This file is part of the "AES tools" project.
# For details, see https://github.com/egor-tensin/aes-tools.
# Distributed under the MIT License.

from collections.abc import Iterable
from enum import Enum
import logging
import os.path
import subprocess


class Algorithm(Enum):
    @staticmethod
    def parse(s):
        return Algorithm(s.lower())

    @staticmethod
    def try_parse(s):
        try:
            return Algorithm.parse(s)
        except ValueError:
            return None

    AES128, AES192, AES256 = "aes128", "aes192", "aes256"

    def __str__(self):
        return self.value


class Mode(Enum):
    @staticmethod
    def parse(s):
        s = s.lower()
        if "{}128".format(Mode.CFB) == s:
            return Mode.CFB
        return Mode(s)

    @staticmethod
    def try_parse(s):
        try:
            return Mode.parse(s)
        except ValueError:
            return None

    ECB, CBC, CFB, OFB, CTR = "ecb", "cbc", "cfb", "ofb", "ctr"

    def requires_init_vector(self):
        return self != Mode.ECB

    def __str__(self):
        return self.value


class BlockInput:
    def __init__(self, key, plaintexts, iv=None):
        self.key = key
        self.plaintexts = plaintexts
        self.iv = iv

    def to_args(self):
        args = [self.key]
        if self.iv is not None:
            args.append(self.iv)
        args.extend(self.plaintexts)
        return args


class Tools:
    def __init__(self, search_dirs, use_sde=False):
        if search_dirs:
            if isinstance(search_dirs, str):
                os.environ["PATH"] += os.pathsep + search_dirs
            elif isinstance(search_dirs, Iterable):
                os.environ["PATH"] += os.pathsep + os.pathsep.join(search_dirs)
            else:
                os.environ["PATH"] += os.pathsep + str(search_dirs)
        self._use_sde = use_sde

    _ENCRYPT_BLOCK = "encrypt_block"
    _DECRYPT_BLOCK = "decrypt_block"
    _ENCRYPT_FILE = "encrypt_file"
    _DECRYPT_FILE = "decrypt_file"

    def run(self, tool_path, args):
        cmd_list = ["sde", "--", tool_path] if self._use_sde else [tool_path]
        cmd_list.extend(args)
        logging.debug("Trying to execute: %s", subprocess.list2cmdline(cmd_list))
        try:
            output = subprocess.check_output(
                cmd_list, universal_newlines=True, stderr=subprocess.STDOUT
            )
        except subprocess.CalledProcessError as e:
            logging.error("Output:\n%s", e.output)
            raise
        logging.debug("Output:\n%s", output)
        return output.split()

    @staticmethod
    def _block_inputs_to_args(inputs):
        args = []
        while True:
            head = next(inputs, None)
            if head is None:
                break
            args.append("--")
            args.extend(head.to_args())
        return args

    @staticmethod
    def _block_settings_to_args(algorithm, mode, use_boxes=False):
        args = [
            "--algorithm",
            str(algorithm),
            "--mode",
            str(mode),
        ]
        if use_boxes:
            args.append("--use-boxes")
        return args

    @staticmethod
    def _build_block_args(algorithm, mode, inputs, use_boxes=False):
        args = Tools._block_settings_to_args(algorithm, mode, use_boxes)
        if isinstance(inputs, Iterable):
            args.extend(Tools._block_inputs_to_args(iter(inputs)))
        else:
            args.extend(inputs.to_args())
        return args

    def run_encrypt_block(self, algorithm, mode, inputs, use_boxes=False):
        args = self._build_block_args(algorithm, mode, inputs, use_boxes)
        return self.run(self._ENCRYPT_BLOCK, args)

    def run_decrypt_block(self, algorithm, mode, inputs, use_boxes=False):
        args = self._build_block_args(algorithm, mode, inputs, use_boxes)
        return self.run(self._DECRYPT_BLOCK, args)

    @staticmethod
    def _build_file_args(algorithm, mode, key, input_path, output_path, iv=None):

        args = [
            "--algorithm",
            str(algorithm),
            "--mode",
            str(mode),
            "--key",
            key,
            "--input",
            input_path,
            "--output",
            output_path,
        ]
        if iv is not None:
            args.extend(("--iv", iv))
        return args

    def run_encrypt_file(self, algorithm, mode, key, input_path, output_path, iv=None):

        args = self._build_file_args(algorithm, mode, key, input_path, output_path, iv)
        return self.run(self._ENCRYPT_FILE, args)

    def run_decrypt_file(self, algorithm, mode, key, input_path, output_path, iv=None):

        args = self._build_file_args(algorithm, mode, key, input_path, output_path, iv)
        return self.run(self._DECRYPT_FILE, args)