# Copyright (c) 2023 Yubico AB
# All rights reserved.
#
#   Redistribution and use in source and binary forms, with or
#   without modification, are permitted provided that the following
#   conditions are met:
#
#    1. Redistributions of source code must retain the above copyright
#       notice, this list of conditions and the following disclaimer.
#    2. Redistributions in binary form must reproduce the above
#       copyright notice, this list of conditions and the following
#       disclaimer in the documentation and/or other materials provided
#       with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

import logging
import os

import click
from cryptography.hazmat.primitives import serialization

from yubikit.core import TRANSPORT
from yubikit.core.smartcard import SW, ApduError, SmartCardConnection
from yubikit.hsmauth import (
    ALGORITHM,
    CREDENTIAL_PASSWORD_LEN,
    MANAGEMENT_KEY_LEN,
    HsmAuthSession,
    InvalidPinError,
)
from yubikit.management import CAPABILITY

from ..hsmauth import (
    generate_random_management_key,
    get_hsmauth_info,
)
from ..util import InvalidPasswordError, parse_private_key
from .util import (
    CliFail,
    click_callback,
    click_force_option,
    click_format_option,
    click_group,
    click_postpone_execution,
    click_prompt,
    get_scp_params,
    log_or_echo,
    pretty_print,
)

logger = logging.getLogger(__name__)


def handle_error(e: Exception, default_exception_msg, target="Credential password"):
    if isinstance(e, InvalidPinError):
        attempts = e.attempts_remaining
        if attempts:
            raise CliFail(f"Wrong {target}, {attempts} attempts remaining.")
        else:
            raise CliFail(f"{target} is blocked.")
    elif isinstance(e, ApduError):
        if e.sw == SW.AUTH_METHOD_BLOCKED:
            raise CliFail("A credential with the provided label already exists.")
        elif e.sw == SW.NO_SPACE:
            raise CliFail("No space left on the YubiKey for YubiHSM Auth credentials.")
        elif e.sw == SW.FILE_NOT_FOUND:
            raise CliFail("Credential with the provided label was not found.")
        elif e.sw == SW.SECURITY_CONDITION_NOT_SATISFIED:
            raise CliFail("The device was not touched.")
        elif e.sw == SW.CONDITIONS_NOT_SATISFIED:
            raise CliFail(f"{target} does not meet complexity requirement.")
    elif isinstance(e, ValueError):
        raise  # Re-raise, ValueErrors are handled elsewhere
    raise CliFail(default_exception_msg)


def _parse_touch_required(touch_required: bool) -> str:
    if touch_required:
        return "On"
    else:
        return "Off"


def _parse_algorithm(algorithm: ALGORITHM) -> str:
    if algorithm == ALGORITHM.AES128_YUBICO_AUTHENTICATION:
        return "Symmetric"
    else:
        return "Asymmetric"


def _parse_hex(hex):
    try:
        return bytes.fromhex(hex)
    except Exception as e:
        raise ValueError(e)


def _parse_key(key, key_len, key_type):
    key = _parse_hex(key)

    if len(key) != key_len:
        raise ValueError(
            f"must be exactly {key_len} bytes long ({key_len * 2} hexadecimal digits) "
            "long"
        )
    return key


def _parse_password(value, key_len, name):
    encoded = value.encode()
    if len(encoded) <= key_len:
        return encoded.ljust(key_len, b"\0")
    if len(encoded) == key_len * 2:
        return _parse_hex(value)
    raise ValueError(f"{name} must be at most 16 bytes")


@click_callback()
def click_parse_management_key(ctx, param, val):
    return _parse_key(val, MANAGEMENT_KEY_LEN, "Management key")


@click_callback()
def click_parse_management_password(ctx, param, val):
    return _parse_password(val, MANAGEMENT_KEY_LEN, "Management password")


@click_callback()
def click_parse_credential_password(ctx, param, val):
    return _parse_password(val, CREDENTIAL_PASSWORD_LEN, "Credential password")


@click_callback()
def click_parse_enc_key(ctx, param, val):
    return _parse_key(val, ALGORITHM.AES128_YUBICO_AUTHENTICATION.key_len, "ENC key")


@click_callback()
def click_parse_mac_key(ctx, param, val):
    return _parse_key(val, ALGORITHM.AES128_YUBICO_AUTHENTICATION.key_len, "MAC key")


@click_callback()
def click_parse_card_crypto(ctx, param, val):
    return _parse_hex(val)


@click_callback()
def click_parse_context(ctx, param, val):
    return _parse_hex(val)


def _prompt_management_key(prompt="Enter management password", confirm=False):
    management_password = click_prompt(
        prompt,
        default="",
        hide_input=True,
        show_default=False,
        confirmation_prompt=confirm,
    )
    return _parse_password(
        management_password, MANAGEMENT_KEY_LEN, "Management password"
    )


def _prompt_credential_password(prompt="Enter credential password", confirm=False):
    credential_password = click_prompt(
        prompt,
        hide_input=True,
        confirmation_prompt=confirm,
    )

    return _parse_password(
        credential_password, CREDENTIAL_PASSWORD_LEN, "Credential password"
    )


def _prompt_symmetric_key(name):
    symmetric_key = click_prompt(f"Enter {name}")

    return _parse_key(
        symmetric_key, ALGORITHM.AES128_YUBICO_AUTHENTICATION.key_len, name
    )


def _fname(fobj):
    return getattr(fobj, "name", fobj)


click_credential_password_option = click.option(
    "-c",
    "--credential-password",
    help="password to protect credential",
    callback=click_parse_credential_password,
)

click_management_key_option = click.option(
    "-m",
    "--management-password",
    "--management-key",
    "management_key",
    help="the management password",
    callback=click_parse_management_password,
)


click_touch_option = click.option(
    "-t", "--touch", is_flag=True, help="require touch on YubiKey to access credential"
)


@click_group(connections=[SmartCardConnection])
@click.pass_context
@click_postpone_execution
def hsmauth(ctx):
    """
    Manage the YubiHSM Auth application
    """
    dev = ctx.obj["device"]
    conn = dev.open_connection(SmartCardConnection)
    ctx.call_on_close(conn.close)

    scp_params = get_scp_params(ctx, CAPABILITY.HSMAUTH, conn)
    ctx.obj["session"] = HsmAuthSession(conn, scp_params)

    try:
        if not scp_params and dev.transport == TRANSPORT.NFC:
            # Dummy command to test access
            ctx.obj["session"].get_management_key_retries()
    except ApduError as e:
        if e.sw == SW.CONDITIONS_NOT_SATISFIED:
            raise CliFail("Unable to manage HSMAuth over NFC without SCP.")
        raise

    info = ctx.obj["info"]
    ctx.obj["fips_unready"] = (
        CAPABILITY.HSMAUTH in info.fips_capable
        and CAPABILITY.HSMAUTH not in info.fips_approved
    )


@hsmauth.command()
@click.pass_context
def info(ctx):
    """
    Display general status of the YubiHSM Auth application.
    """
    info = ctx.obj["info"]
    data = get_hsmauth_info(ctx.obj["session"])
    if CAPABILITY.HSMAUTH in info.fips_capable:
        # This is a bit ugly as it makes assumptions about the structure of data
        data["FIPS approved"] = CAPABILITY.HSMAUTH in info.fips_approved

    click.echo("\n".join(pretty_print(data)))


@hsmauth.command()
@click.pass_context
@click_force_option
def reset(ctx, force):
    """
    Reset all YubiHSM Auth data.

    This action will wipe all data and restore factory setting for
    the YubiHSM Auth application on the YubiKey.
    """

    if not force:
        click.confirm(
            "WARNING! This will delete all stored YubiHSM Auth data and restore "
            "factory setting. Proceed?",
            abort=True,
            err=True,
        )

    click.echo("Resetting YubiHSM Auth data...")
    ctx.obj["session"].reset()

    click.echo(
        "Reset complete. All YubiHSM Auth data has been cleared from the YubiKey."
    )
    click.echo("Your YubiKey now has an empty Management password.")


@hsmauth.group()
def credentials():
    """Manage YubiHSM Auth credentials."""


@credentials.command()
@click.pass_context
def list(ctx):
    """
    List all credentials.

    List all credentials stored on the YubiKey.
    """
    session = ctx.obj["session"]
    creds = session.list_credentials()

    if len(creds) == 0:
        click.echo("No items found")
    else:
        click.echo(f"Found {len(creds)} item(s)")

        max_size_label = max(len(cred.label) for cred in creds)
        max_size_type = (
            10
            if any(
                c.algorithm == ALGORITHM.EC_P256_YUBICO_AUTHENTICATION for c in creds
            )
            else 9
        )

        format_str = "{0: <{label_width}}\t{1: <{type_width}}\t{2}\t{3}"

        click.echo(
            format_str.format(
                "Label",
                "Type",
                "Touch",
                "Retries",
                label_width=max_size_label,
                type_width=max_size_type,
            )
        )

        for cred in creds:
            click.echo(
                format_str.format(
                    cred.label,
                    _parse_algorithm(cred.algorithm),
                    _parse_touch_required(cred.touch_required),
                    cred.counter,
                    label_width=max_size_label,
                    type_width=max_size_type,
                )
            )


@credentials.command()
@click.pass_context
@click.argument("label")
@click_credential_password_option
@click_management_key_option
@click_touch_option
def generate(ctx, label, credential_password, management_key, touch):
    """Generate an asymmetric credential.

    This will generate an asymmetric YubiHSM Auth credential
    (private key) on the YubiKey.

    \b
    LABEL label for the YubiHSM Auth credential
    """

    if ctx.obj["fips_unready"]:
        raise CliFail(
            "YubiKey FIPS must be in FIPS approved mode prior to adding credentials."
        )

    if management_key is None:
        management_key = _prompt_management_key()

    if credential_password is None:
        credential_password = _prompt_credential_password()

    session = ctx.obj["session"]

    try:
        session.generate_credential_asymmetric(
            management_key, label, credential_password, touch
        )
        click.echo("Asymmetric credential generated.")
    except Exception as e:
        handle_error(
            e, default_exception_msg="Failed to generate asymmetric credential."
        )


@credentials.command("import")
@click.pass_context
@click.argument("label")
@click.argument("private-key", type=click.File("rb"), metavar="PRIVATE-KEY")
@click.option("-p", "--password", help="password used to decrypt the private key")
@click_credential_password_option
@click_management_key_option
@click_touch_option
def import_credential(
    ctx, label, private_key, password, credential_password, management_key, touch
):
    """Import an asymmetric credential.

    This will import a private key as an asymmetric YubiHSM Auth credential
    to the YubiKey.

    \b
    LABEL        label for the YubiHSM Auth credential
    PRIVATE-KEY  file containing the private key (use '-' to use stdin)
    """
    if ctx.obj["fips_unready"]:
        raise CliFail(
            "YubiKey FIPS must be in FIPS approved mode prior to adding credentials."
        )

    if management_key is None:
        management_key = _prompt_management_key()

    if credential_password is None:
        credential_password = _prompt_credential_password()

    session = ctx.obj["session"]

    data = private_key.read()

    while True:
        if password is not None:
            password = password.encode()
        try:
            private_key = parse_private_key(data, password)
        except InvalidPasswordError:
            logger.debug("Error parsing key", exc_info=True)
            if password is None:
                password = click_prompt(
                    "Enter password to decrypt key",
                    default="",
                    hide_input=True,
                    show_default=False,
                )
                continue
            else:
                password = None
                click.echo("Wrong password.")
            continue
        break

    try:
        session.put_credential_asymmetric(
            management_key,
            label,
            private_key,
            credential_password,
            touch,
        )
        click.echo("Asymmetric credential imported.")
    except Exception as e:
        handle_error(e, default_exception_msg="Failed to import asymmetric credential.")


@credentials.command()
@click.pass_context
@click.argument("label")
@click.argument("public-key-output", type=click.File("wb"), metavar="PUBLIC-KEY")
@click_format_option
def export(ctx, label, public_key_output, format):
    """Export the public key corresponding to an asymmetric credential.

    This will export the long-term public key corresponding to the
    asymmetric YubiHSM Auth credential stored on the YubiKey.

    \b
    LABEL      label for the YubiHSM Auth credential
    PUBLIC-KEY file to write the public key to (use '-' to use stdout)
    """

    session = ctx.obj["session"]

    try:
        public_key = session.get_public_key(label)
        key_encoding = format
        public_key_encoded = public_key.public_bytes(
            encoding=key_encoding,
            format=serialization.PublicFormat.SubjectPublicKeyInfo,
        )

        public_key_output.write(public_key_encoded)

        log_or_echo(
            f"Public key for {label} written to {_fname(public_key_output)}",
            logger,
            public_key_output,
        )
    except ApduError as e:
        if e.sw == SW.AUTH_METHOD_BLOCKED:
            raise CliFail("The entry is not an asymmetric credential.")
        elif e.sw == SW.FILE_NOT_FOUND:
            raise CliFail("Credential not found.")
        else:
            raise CliFail("Unable to export public key.")


@credentials.command()
@click.pass_context
@click.argument("label")
@click.option("-E", "--enc-key", help="the ENC key", callback=click_parse_enc_key)
@click.option("-M", "--mac-key", help="the MAC key", callback=click_parse_mac_key)
@click.option(
    "-g", "--generate", is_flag=True, help="generate a random encryption and mac key"
)
@click_credential_password_option
@click_management_key_option
@click_touch_option
def symmetric(
    ctx, label, credential_password, management_key, enc_key, mac_key, generate, touch
):
    """Import a symmetric credential.

    This will import an encryption and mac key as a symmetric YubiHSM Auth credential on
    the YubiKey.

    \b
    LABEL  label for the YubiHSM Auth credential
    """

    if ctx.obj["fips_unready"]:
        raise CliFail(
            "YubiKey FIPS must be in FIPS approved mode prior to adding credentials."
        )

    if management_key is None:
        management_key = _prompt_management_key()

    if credential_password is None:
        credential_password = _prompt_credential_password()

    if generate and (enc_key or mac_key):
        ctx.fail("--enc-key and --mac-key cannot be combined with --generate")

    if generate:
        enc_key = os.urandom(ALGORITHM.AES128_YUBICO_AUTHENTICATION.key_len)
        mac_key = os.urandom(ALGORITHM.AES128_YUBICO_AUTHENTICATION.key_len)
        click.echo("Generated ENC and MAC keys:")
        click.echo("\n".join(pretty_print({"ENC-KEY": enc_key, "MAC-KEY": mac_key})))

    if not enc_key:
        enc_key = _prompt_symmetric_key("ENC key")

    if not mac_key:
        mac_key = _prompt_symmetric_key("MAC key")

    session = ctx.obj["session"]

    try:
        session.put_credential_symmetric(
            management_key,
            label,
            enc_key,
            mac_key,
            credential_password,
            touch,
        )
        click.echo("Symmetric credential stored.")

    except Exception as e:
        handle_error(e, default_exception_msg="Failed to import symmetric credential.")


@credentials.command()
@click.pass_context
@click.argument("label")
@click.option(
    "-d", "--derivation-password", help="deriviation password for ENC and MAC keys"
)
@click_credential_password_option
@click_management_key_option
@click_touch_option
def derive(ctx, label, derivation_password, credential_password, management_key, touch):
    """Import a symmetric credential derived from a password.

    This will import a symmetric YubiHSM Auth credential by deriving
    ENC and MAC keys from a password.

    \b
    LABEL  label for the YubiHSM Auth credential
    """

    if ctx.obj["fips_unready"]:
        raise CliFail(
            "YubiKey FIPS must be in FIPS approved mode prior to adding credentials."
        )

    if management_key is None:
        management_key = _prompt_management_key()

    if credential_password is None:
        credential_password = _prompt_credential_password()

    if derivation_password is None:
        derivation_password = click_prompt(
            "Enter derivation password",
            hide_input=True,
            confirmation_prompt=True,
        )

    session = ctx.obj["session"]

    try:
        session.put_credential_derived(
            management_key, label, derivation_password, credential_password, touch
        )
        click.echo("Derived symmetric credential stored.")
    except Exception as e:
        handle_error(e, default_exception_msg="Failed to import symmetric credential.")


@credentials.command()
@click.pass_context
@click.argument("label")
@click_management_key_option
@click_force_option
def delete(ctx, label, management_key, force):
    """
    Delete a credential.

    This will delete a YubiHSM Auth credential from the YubiKey.

    \b
    LABEL a label to match a single credential (as shown in "list")
    """

    if management_key is None:
        management_key = _prompt_management_key()

    if not force:
        click.confirm(
            f"Delete credential: {label} ?",
            abort=True,
            err=True,
        )

    session = ctx.obj["session"]

    try:
        session.delete_credential(management_key, label)
        click.echo("Credential deleted.")
    except Exception as e:
        handle_error(
            e,
            default_exception_msg="Failed to delete credential.",
        )


@credentials.command()
@click.pass_context
@click.argument("label")
@click.option(
    "-c",
    "--credential-password",
    help="current credential password",
    callback=click_parse_credential_password,
)
@click_management_key_option
@click.option(
    "-n",
    "--new-credential-password",
    help="a new credential password to set",
    callback=click_parse_credential_password,
)
def change_password(
    ctx, label, credential_password, management_key, new_credential_password
):
    """
    Change the password of a credential.

    This will change the password of a YubiHSM Auth credential stored on the YubiKey.

    Changing a password requires either the current password for the credential,
    or the management password (not both).

    \b
    LABEL a label to match a single credential (as shown in "list")
    """

    session = ctx.obj["session"]
    if credential_password and management_key:
        ctx.fail("--credential-password and --management-password can't be combined")

    if management_key is None and credential_password is None:
        click.echo("Choose authentication method:")
        click.echo(" [1] Current credential password")
        click.echo(" [2] Management password")
        method = click_prompt(
            "Enter selection",
            type=click.Choice([1, 2]),
            default=1,
            show_choices=False,
        )
        if method == 1:
            credential_password = _prompt_credential_password(
                "Enter current credential password"
            )
        else:
            management_key = _prompt_management_key()

    use_management_key = management_key is not None
    if new_credential_password is None:
        new_credential_password = _prompt_credential_password(
            "Enter the new credential password", confirm=True
        )

    try:
        if use_management_key:
            session.change_credential_password_admin(
                label, management_key, new_credential_password
            )
        else:
            session.change_credential_password(
                label, credential_password, new_credential_password
            )
        click.echo("Credential password changed.")
    except Exception as e:
        handle_error(
            e,
            default_exception_msg="Failed to change credential password.",
            target="Management password"
            if use_management_key
            else "Credential password",
        )


@hsmauth.group()
def access():
    """Manage Management Key for YubiHSM Auth"""


@access.command(hidden=True)
@click.pass_context
@click.option(
    "-m",
    "--management-key",
    help="current management key",
    show_default=True,
    callback=click_parse_management_key,
)
@click.option(
    "-n",
    "--new-management-key",
    help="a new management key to set",
    callback=click_parse_management_key,
)
@click.option(
    "-g",
    "--generate",
    is_flag=True,
    help="generate a random management key (can't be used with --new-management-key)",
)
def change_management_key(ctx, management_key, new_management_key, generate):
    """
    Change the management key.

    Allows you to change the management key which is required to add and delete
    YubiHSM Auth credentials stored on the YubiKey.
    """

    if management_key is None:
        management_key = _prompt_management_key(
            "Enter current management key [blank to use default key]"
        )

    session = ctx.obj["session"]

    # Can't combine new key with generate.
    if new_management_key and generate:
        ctx.fail("Invalid options: --new-management-key conflicts with --generate")

    if new_management_key is None:
        if generate:
            if ctx.obj["info"].version >= (5, 8):
                raise CliFail(
                    "Generating a random management key is not supported on this "
                    "YubiKey."
                )
            new_management_key = generate_random_management_key()
            click.echo(f"Generated management key: {new_management_key.hex()}")
        else:
            try:
                new_management_key = bytes.fromhex(
                    click_prompt(
                        "Enter the new management key",
                        hide_input=True,
                        confirmation_prompt=True,
                    )
                )
            except Exception:
                ctx.fail("New management key has the wrong format.")

    if len(new_management_key) != MANAGEMENT_KEY_LEN:
        raise CliFail(
            "Management key has the wrong length (expected %d bytes)."
            % MANAGEMENT_KEY_LEN
        )

    try:
        session.put_management_key(management_key, new_management_key)
        click.echo("Management key changed.")
    except Exception as e:
        handle_error(
            e,
            default_exception_msg="Failed to change management key.",
            target="Management key",
        )


@access.command()
@click.pass_context
@click.option(
    "-m",
    "--management-password",
    "management_key",
    help="current management password",
    show_default=True,
    callback=click_parse_management_password,
)
@click.option(
    "-n",
    "--new-management-password",
    "new_management_key",
    help="a new management password to set",
    callback=click_parse_management_password,
)
def change_management_password(ctx, management_key, new_management_key):
    """
    Change the management key.

    Allows you to change the management key which is required to add and delete
    YubiHSM Auth credentials stored on the YubiKey.
    """

    if management_key is None:
        management_key = _prompt_management_key(
            "Enter your current management password",
        )

    if new_management_key is None:
        new_management_key = _prompt_management_key(
            "Enter a new management password", confirm=True
        )

    session = ctx.obj["session"]
    try:
        session.put_management_key(management_key, new_management_key)
        click.echo("Management password changed.")
    except Exception as e:
        handle_error(
            e,
            default_exception_msg="Failed to change management password.",
            target="Management password",
        )