"""Lambda aws_break_glass function module."""
import base64
import datetime
import json
from collections import OrderedDict

import boto3
import botocore
from botocore.config import Config
from lambdacommon.common_config import logger
import sentry_sdk
from sentry_sdk.integrations.aws_lambda import AwsLambdaIntegration

import config
from src.utils import jira
from src.utils import notifications
from src.utils import sqs

sentry_sdk.init(
    dsn=config.SENTRY_DSN,
    integrations=[AwsLambdaIntegration()],
    traces_sample_rate=1.0,
)

iam_client_config = Config(
    retries={
        'total_max_attempts': 10,
        'mode': 'adaptive'
    }
)

client = boto3.client('iam', config=iam_client_config)
sts_client = boto3.client('sts')


def create_and_attach_policy(user, policy_document):
    """
    Create IAM policy and attach to user.

    Args:
        user (str): Name of user
        policy_document (json): Formatted IAM policy document
    Returns:
        str: Resulting policy name
    """
    now = datetime.datetime.now()
    now_formatted = now.strftime('%Y-%m-%d-%H-%M-%S')
    try:
        response = client.create_policy(
            PolicyName=f'BreakGlass-{user}-{now_formatted}',
            PolicyDocument=policy_document,
            Description='Temporary policy to assume break glass role',
            Tags=[
                {
                    'Key': 'service_name',
                    'Value': 'break_glass'
                },
                {
                    'Key': 'temporary',
                    'Value': 'true'
                },
            ]
        )
        logger.info('New policy ARN: {}'.format(response['Policy']['Arn']))
        policy_attachment = client.attach_user_policy(
            UserName=user,
            PolicyArn=response['Policy']['Arn']
        )
        # The response object is effectively empty so check the status code
        if policy_attachment['ResponseMetadata']['HTTPStatusCode'] == 200:
            logger.info(f'Policy successfully attached to {user}')
    except botocore.exceptions.ClientError as error:
        logger.exception('Error creating or attaching IAM policy: {}'.format(
            error.response['Error']['Message']))
        raise error
    return response['Policy']['PolicyName']


def format_policy_document(json_template_file, break_glass_role_arn):  # noqa: E501
    """
    Format an IAM policy document from template.

    Args:
        json_template_file (str): Name of policy template file to open
        break_glass_role_arn (str): ARN of break glass role
    Returns:
        json: Formatted policy document
    """
    policy_file = open(json_template_file, 'r')
    policy = json.load(policy_file, object_pairs_hook=OrderedDict)
    assume_expire_at = datetime.datetime.now() + \
        datetime.timedelta(minutes=config.ASSUME_ROLE_TIMEOUT_MINUTES)
    assume_expire_at_fmt = assume_expire_at.strftime('%Y-%m-%dT%H:%M:%SZ')
    policy['Statement'][0]['Condition'][
        'DateLessThan']['aws:CurrentTime'] = assume_expire_at_fmt
    policy['Statement'][0]['Resource'] = break_glass_role_arn

    return policy


def get_aws_user_by_email_tag(email: str) -> str:
    """
    Find aws user by value of email tag.

    :param email: email tag value to search for
    :type email: str
    :return: username
    :rtype: str
    """
    paginator = client.get_paginator("list_users")
    search_results = []
    for page in paginator.paginate():
        for user in page["Users"]:
            username = user["UserName"]
            tags = client.list_user_tags(UserName=username)["Tags"]
            for tag in tags:
                if tag["Key"] == "email" and tag["Value"] == email:
                    search_results.append(username)
    if len(search_results) == 1:
        return search_results[0]
    elif len(search_results) == 0:
        logger.error(f"Cannot find user with email tag {email}")
        raise ValueError(f"Cannot find user with email tag {email}")
    else:
        users_str = ','.join(search_results)
        logger.error(
            f"Found multiple users with email tag {email}: {users_str}")
        raise ValueError(
            f"Found multiple users with email tag {email}: {users_str}")


def get_user_from_event(event):
    """
    Parse event and return AWS user.

    Args:
        event (json): AWS Lambda event
    Returns:
        str: An IAM user derived from OIDC email address
    """
    oidc_jwt = event['headers']['x-amzn-oidc-data']
    # Get the JWT payload
    base64_payload = oidc_jwt.split('.')[1]
    byte_payload = base64.b64decode(base64_payload)
    payload = byte_payload.decode('utf-8')
    json_payload = json.loads(payload)
    # Payload is different in Azure and Okta requests
    # In Okta email field is 'preferred_username'
    # In Azure it is 'email'
    return get_aws_user_by_email_tag(json_payload['email'])


def generate_break_glass_command(user, mfa_serial_number, break_glass_role_arn):  # noqa: E501
    """
    Generate a break glass command.

    Args:
        user (str): Name of user
        mfa_serial_number (str): Serial number of MFA device
        break_glass_role_arn (str): ARN of break glass role
    Returns:
        str: A command meant to be run by a user
    """
    command = f'eval $(aws sts assume-role --role-arn {break_glass_role_arn} '\
        f'--role-session-name {user} --serial-number {mfa_serial_number} '\
        "--token-code MFA_TOKEN_CODE | jq -r '.Credentials | "\
        '"export AWS_ACCESS_KEY_ID=\(.AccessKeyId)\\nexport '\
        'AWS_SECRET_ACCESS_KEY=\(.SecretAccessKey)\\nexport '\
        'AWS_SESSION_TOKEN=\(.SessionToken)\\n"\')'  # noqa: W605
    explanation = 'Run this command within 15 minutes, substituting '\
        'a real MFA code for MFA_TOKEN_CODE'
    return f'{explanation}: {command}'


def verify_break_glass_role(role_name, verification_role_arn):
    """
    Look up break glass role and verify it exists.

    Args:
        role_name (str): Name of role to look up
        verification_role_arn (str): ARN of verification role
    Returns:
        str: Role ARN
    """
    try:
        assumed_role = sts_client.assume_role(
            RoleArn=verification_role_arn,
            RoleSessionName=config.ROLE_SESSION_NAME
        )
        ACCESS_KEY = assumed_role['Credentials']['AccessKeyId']
        SECRET_KEY = assumed_role['Credentials']['SecretAccessKey']
        SESSION_TOKEN = assumed_role['Credentials']['SessionToken']
        iam_client = boto3.client(
            'iam',
            aws_access_key_id=ACCESS_KEY,
            aws_secret_access_key=SECRET_KEY,
            aws_session_token=SESSION_TOKEN,
        )
        break_glass_role = iam_client.get_role(RoleName=role_name)
        return break_glass_role['Role']['Arn']
    except botocore.exceptions.ClientError as error:
        logger.exception('Error looking up IAM role: {}'.format(
            error.response['Error']['Message']))
        raise error


def verify_mfa_device(user):
    """
    Look up MFA device for a user and verify it exists.

    Args:
        user (str): Name of user
    Returns:
        str: Serial number of MFA device
    """
    try:
        mfa_devices = client.list_mfa_devices(UserName=user)
        assert mfa_devices['MFADevices'], \
            f'{user} does not have any configured MFA devices'
        return mfa_devices['MFADevices'][0]['SerialNumber']
    except botocore.exceptions.ClientError as error:
        logger.exception('Error looking up MFA devices for {}: {}'.format(
            user, error.response['Error']['Message']))
        raise error


def handler(event, context):
    """Lambda entry point."""
    try:
        user = get_user_from_event(event)

        # Get AWS account id from the request
        account_id = event['queryStringParameters']['account_id']

        # Verify that jira issue exists
        jira_issue = event['queryStringParameters']['jira_issue']
        issue_url = jira.verify_jira_issue(jira_issue)
        logger.info(f'Jira issue {jira_issue} found with url {issue_url}')

        # Verify that break glass role exists
        break_glass_role_name = config.BREAK_GLASS_ROLE_NAME.get(account_id)
        verification_role_arn = config.VERIFICATION_AWS_ROLE_ARN.get(account_id)  # noqa: E501
        break_glass_role_arn = verify_break_glass_role(break_glass_role_name,
                                                       verification_role_arn)

        # Create and attach policy to user
        policy_document = format_policy_document(
            'policy.json', break_glass_role_arn)
        policy_name = create_and_attach_policy(
            user, json.dumps(policy_document))
        logger.info(f'{policy_name} now permits assume role access. '
                    'Generating break glass command...')

        # Verify that user has MFA device
        mfa_serial_number = verify_mfa_device(user)

        # Generate break glass command
        assume_role_command = generate_break_glass_command(
            user, mfa_serial_number, break_glass_role_arn)

        # Add comment to Jira issue
        comment_timestamp = jira.comment_on_jira_issue(jira_issue, user, account_id)  # noqa: E501
        logger.info(f'Added comment to Jira issue {jira_issue} at time '
                    f'{comment_timestamp}')

        # Send message to SQS for downstream auditing
        sqs_message_id = sqs.send_message(jira_issue, user, account_id)
        logger.info(f'Message sent to SQS with ID {sqs_message_id}')

        # Send notifications via email and slack
        message_id = notifications.send_email(
            account_id, issue_url, user)
        logger.info(f'Email notification sent with message ID {message_id}')

        slack_timestamp, slack_channel = notifications.send_slack_message(
            account_id, issue_url, user)
        logger.info(f'Slack notification sent to {slack_channel} '
                    f'and received at {slack_timestamp}')

        console_url = 'https://signin.aws.amazon.com/switchrole?'\
            'roleName={}&account={}'.format(break_glass_role_name, account_id)
        response = {
            'statusCode': 200,
            'statusDescription': '200 OK',
            'isBase64Encoded': False,
            'headers': {
                'Content-Type': 'text/html'
            },
            'body': f'<h1>CLI</h1><pre>{assume_role_command}<pre>'
            f'<h1>Console</h1><a href={console_url}>'
            'Switch to break glass role</a>'
        }
        return response
    except Exception as error:
        logger.exception(str(error))
        sentry_sdk.capture_exception(error)
        response = {
            'statusCode': 500,
            'statusDescription': '500 Internal Server Error',
            'isBase64Encoded': False,
            'headers': {
                'Content-Type': 'text/html'
            },
            'body': '<pre>{}<pre>'.format(str(error))
        }
        return response