"""A collection of helper functions for custom policies."""

import os
import re

import boto3

AWS_SESSION = None


def flatten(data):
    """Flatten a checkov conf parameter list by returning the first element.

    If the list contains more than one element, or the data is not a list
    at all, it gets returned without any changes.

    Args:
        data (mixed): Configuration parameter.

    Returns:
        mixed: Flattened configuration parameter.
    """
    if not isinstance(data, list):
        return data

    if len(data) != 1:
        return data

    return data[0]


def get_module_name(data):
    """Extract <org-name>/<repo-name> part from the `source` attribute.

    Args:
        data (mixed): Module configuration.

    Returns:
        string: Module name or `None`.
    """
    if data and 'source' in data:
        source_attribute = flatten(data.get('source'))
        match = re.search(r'(?:git@|https://)github\.com[:/]'
                          r'([\w\/\-\.]+?)(\.git)?(//)?[/\w]*?'
                          r'\?ref=([\w\-\.]+)', source_attribute)
        if match:
            return match.group(1)

    return None


def get_aws_session():
    """Create or re-use an AWS session.

    If the session has already been established, use it instead of creating
    a new one. If environment variable CHECKOV_AWS_ROLE exists, assume the
    role ARN specified in this env.

    No exception handling happens in this function. It should be the check's
    choice to return failure or skip.

    Returns:
        boto3.session.Session: AWS session.
    """
    global AWS_SESSION

    if AWS_SESSION:
        return AWS_SESSION

    CHECKOV_AWS_ROLE = os.getenv('CHECKOV_AWS_ROLE')
    if not CHECKOV_AWS_ROLE:
        AWS_SESSION = boto3.Session()
        return AWS_SESSION

    client = boto3.client('sts')
    response = client.assume_role(RoleArn=CHECKOV_AWS_ROLE,
                                  RoleSessionName='checkov')
    sts_credentials = {
        'aws_access_key_id': response['Credentials']['AccessKeyId'],
        'aws_secret_access_key': response['Credentials']['SecretAccessKey'],
        'aws_session_token': response['Credentials']['SessionToken']
    }
    AWS_SESSION = boto3.Session(**sts_credentials)

    return AWS_SESSION