"""Lambda create-shareable-snapshot function module."""
import datetime
import json
import logging
import sys

import boto3
import sentry_sdk
from sentry_sdk.integrations.aws_lambda import AwsLambdaIntegration
from tenacity import before_sleep_log
from tenacity import Retrying
from tenacity import retry_if_exception_type
from tenacity import stop_after_attempt
from tenacity import wait_fixed

import config

from common.logic import database # noqa

if config.SENTRY_DSN:
    sentry_sdk.init(
        dsn=config.SENTRY_DSN,
        environment=config.ENVIRONMENT,
        integrations=[AwsLambdaIntegration(timeout_warning=True)]
    )

logger = logging.getLogger()
logger.setLevel(config.LOGGER_LEVEL)
stream_handler = logging.StreamHandler(sys.stdout)
stream_handler.setLevel(logging.DEBUG)
logger.handlers = [stream_handler]


def assume_source_account_role(account_id):
    """
    Assume IAM role in source account.

    Args:
    account_id (str): AWS account ID

    Returns:
    dict: dict of AWS credentials
    """
    client = boto3.client('sts')
    assume_role_response = client.assume_role(
        RoleArn=f'arn:aws:iam::{account_id}:role/{config.CROSS_ACCOUNT_BACKUP_ROLE_NAME}', # noqa
        RoleSessionName=config.SERVICE_NAME,
        ExternalId=config.EXTERNAL_ID,
        DurationSeconds=3600
    )
    credentials = assume_role_response['Credentials']

    return credentials


def create_initial_snapshot(db_name, db_type, rds_client):
    """
    Create initial manual snapshot.

    Args:
    db_name (str): DB name
    db_name (str): DB type
    rds_client (obj): boto3 RDS client

    Returns:
    str: resulting snapshot name
    """
    # Wait for database to be available in case it is currently backing up
    database.wait_for_database_availability(
        db_name, db_type, client=rds_client)

    # Take snapshot
    snapshot_datetime = datetime.datetime.now()
    snapshot_time = snapshot_datetime.strftime('%Y-%m-%d-%H-%M-%S')
    snapshot_name = f'{db_name}-{snapshot_time}'

    logger.info(
        f'Creating snapshot {snapshot_name} of database {db_name}')
    # It's possible for the availability check to succeed and the snapshot
    # operation to still fail if another snapshot operation has been issued
    # recently, because there is a delay between a snapshot operation
    # being issued and the cluster status changing. We therefore retry on the
    # specific exceptions that are produced in this scenario.
    for attempt in Retrying(
            stop=stop_after_attempt(config.SNAPSHOT_RETRY_LIMIT),
            wait=wait_fixed(config.SNAPSHOT_RETRY_DELAY),
            retry=retry_if_exception_type(
                rds_client.exceptions.InvalidDBClusterStateFault
                | rds_client.exceptions.InvalidDBClusterSnapshotStateFault # noqa
            ),
            before_sleep=before_sleep_log(logger, logging.INFO)):
        with attempt:
            snapshot = database.create_snapshot(
                snapshot_name, db_name,
                db_type, client=rds_client, additional_tags=get_snapshot_tags()) # noqa

    return snapshot


def create_shareable_snapshot(
        db_type, rds_client, snapshot_name, source_account_id):
    """
    Create shareable snapshot.

    Args:
    db_type (str): DB type
    rds_client (obj): boto3 RDS client
    snapshot_name (str): initial snapshot name
    source_account_id (str): ID of source AWS account

    Returns:
    str: shareable snapshot name
    """
    # Now create a shareable snapshot from the initial snapshot
    shared_snapshot_name = f'{snapshot_name}-shared'
    logger.info(
        f'Creating shareable snapshot {shared_snapshot_name} from '
        f'{snapshot_name}')

    shareable_snapshot_id = database.copy_snapshot(
        snapshot_name,
        source_account_id,
        db_type,
        shared_snapshot_name,
        'shareable',
        config.SHARE_KMS_KEY_ID,
        config.AWS_DEFAULT_REGION,
        client=rds_client,
        additional_tags=get_snapshot_tags())

    return shareable_snapshot_id


def get_target_account_id():
    """
    Get the account ID to share the snapshot with.

    If not explicitly defined, defaults to the current account ID.

    Returns:
    str: AWS account ID to share the snapshot with.
    """
    if config.TARGET_ACCOUNT_ID:
        return config.TARGET_ACCOUNT_ID

    sts_client = boto3.client('sts')
    target_account_id = sts_client.get_caller_identity()['Account']

    return target_account_id


def get_snapshot_tags():
    """Return any additional tags to add to snapshots."""
    return [
        {'Key': 'execution_id', 'Value': config.EXECUTION_ID},
    ]


def encrypted_with_default_key(db_name, db_type, rds_client, kms_client):
    """Return true if the given database is encrypted with the default key."""
    db_config = database.get_config(db_name, db_type, rds_client)
    kms_key_arn = db_config.get('KmsKeyId')

    default_key = kms_client.describe_key(
        KeyId='alias/aws/rds'
    )

    return kms_key_arn == default_key['KeyMetadata']['Arn']


def snapshot_encrypted_with_default_key(
        snapshot_name, db_type, rds_client, kms_client):
    """Return true if the given snapshot is encrypted with the default key."""
    if db_type.lower() == 'cluster':
        response = rds_client.describe_db_cluster_snapshots(
            DBClusterSnapshotIdentifier=snapshot_name)
        kms_key_arn = response['DBClusterSnapshots'][0].get('KmsKeyId')
    else:
        response = rds_client.describe_db_snapshots(
            DBSnapshotIdentifier=snapshot_name)
        kms_key_arn = response['DBSnapshots'][0].get('KmsKeyId')

    default_key = kms_client.describe_key(
        KeyId='alias/aws/rds'
    )

    return kms_key_arn == default_key['KeyMetadata']['Arn']


def main():
    """Task entry point."""
    sfn_client = boto3.client(
        'stepfunctions', region_name=config.AWS_DEFAULT_REGION)
    try:
        target_account_id = get_target_account_id()
        credentials = assume_source_account_role(config.SOURCE_ACCOUNT_ID)
        rds_client = boto3.client(
            'rds',
            region_name=config.AWS_DEFAULT_REGION,
            aws_access_key_id=credentials['AccessKeyId'],
            aws_secret_access_key=credentials['SecretAccessKey'],
            aws_session_token=credentials['SessionToken']
        )

        if config.RESTORE_SNAPSHOT_ID:
            logger.info(
                f'Using supplied snapshot {config.RESTORE_SNAPSHOT_ID} '
                f'instead of creating a new one. The supplied snapshot '
                f'lifecycle is owned by the requesting team.')
            snapshot_name = config.RESTORE_SNAPSHOT_ID
        else:
            snapshot_name = create_initial_snapshot(
                config.DB_NAME, config.DB_TYPE, rds_client)

        # If the target account is different to the source account, we need to
        # share the snapshot with the target account.
        if config.SOURCE_ACCOUNT_ID != target_account_id:
            kms_client = boto3.client(
                'kms',
                region_name=config.AWS_DEFAULT_REGION,
                aws_access_key_id=credentials['AccessKeyId'],
                aws_secret_access_key=credentials['SecretAccessKey'],
                aws_session_token=credentials['SessionToken']
            )
            # If the snapshot is encrypted with the default KMS key, we need
            # to create an intermediate snapshot that is encrypted with a
            # CMK in order to be able to share it. A supplied snapshot may use
            # a different key to the current database, so inspect the snapshot
            # itself rather than the live database configuration.
            if config.RESTORE_SNAPSHOT_ID:
                is_default_key = snapshot_encrypted_with_default_key(
                    snapshot_name, config.DB_TYPE, rds_client, kms_client)
            else:
                is_default_key = encrypted_with_default_key(
                    config.DB_NAME, config.DB_TYPE, rds_client, kms_client)

            if is_default_key:
                snapshot_name = create_shareable_snapshot(
                    config.DB_TYPE, rds_client,
                    snapshot_name, config.SOURCE_ACCOUNT_ID)

            logger.info(
                f'Configuring sharing settings for {snapshot_name}')
            database.configure_snapshot_sharing_settings(
                snapshot_name, target_account_id,
                config.DB_TYPE, client=rds_client)

        if config.TASK_TOKEN:
            sfn_client.send_task_success(
                taskToken=config.TASK_TOKEN,
                output=json.dumps(snapshot_name)
            )
        return snapshot_name
    except Exception as error:
        logger.exception(str(error))
        if config.TASK_TOKEN:
            sfn_client.send_task_failure(
                taskToken=config.TASK_TOKEN,
                error='500',
                cause=str(error),
            )
        raise error


if __name__ == '__main__':
    main()