import json
import logging
# from botocore.vendored import requests
import random
import string
import time
import boto3
import os
import requests
from botocore.exceptions import ClientError

# Falcon SDK - Cloud_Connect_AWS and OAuth2 API service classes
from falconpy import cspm_registration as CSPM
from falconpy import oauth2 as FalconAuth

logger = logging.getLogger()
logger.setLevel(logging.INFO)

CF = "https://cs-prod-cloudconnect-templates.s3.amazonaws.com/aws_cspm_cloudformation_v2.json"
CF_NO_CT = "https://cs-prod-cloudconnect-templates.s3.amazonaws.com/aws_cspm_cloudformation_v2_no_cloudtrail.json"
STACKSETURL = "https://cs-prod-cloudconnect-templates.s3.amazonaws.com/aws_cspm_cloudformation_iam_v2.json"
STACKSETNAME = 'CrowdStrike-CSPM-Integration'
STACKNAME = 'CrowdStrike-CSPM-Integration'

SUCCESS = "SUCCESS"
FAILED = "FAILED"
CT_REGION = os.environ['CT_REGION']
SECRET_NAME = os.environ['SECRET_ARN']
QSS3BucketName = os.environ['QSS3BucketName']
CSRoleName = os.environ['CSRoleName']
CSAccountNumber = os.environ['CSAccountNumber']


def cfnresponse_send(event, context, responseStatus, responseData, physicalResourceId=None, noEcho=False):
    responseUrl = event['ResponseURL']
    print(responseUrl)

    responseBody = {'Status': responseStatus,
                    'Reason': 'See the details in CloudWatch Log Stream: ' + context.log_stream_name,
                    'PhysicalResourceId': physicalResourceId or context.log_stream_name, 'StackId': event['StackId'],
                    'RequestId': event['RequestId'], 'LogicalResourceId': event['LogicalResourceId'], 'NoEcho': noEcho,
                    'Data': responseData}

    json_responseBody = json.dumps(responseBody)

    print("Response body:\n" + json_responseBody)

    headers = {
        'content-type': '',
        'content-length': str(len(json_responseBody))
    }

    try:
        response = requests.put(responseUrl,
                                data=json_responseBody,
                                headers=headers)
        print("Status code: " + response.reason)
    except Exception as e:
        print("send(..) failed executing requests.put(..): " + str(e))


def get_secret_value(secret):
    # key = 'falcon_client_secret'
    key = secret
    SM = boto3.client('secretsmanager')
    secret_list = SM.list_secrets()['SecretList']
    output = {}
    for s in secret_list:
        if key in s.values():
            output = SM.get_secret_value(SecretId=key)['SecretString']
    return output


def get_master_id():
    """ Get the master Id from AWS Organization - Only on master"""
    masterID = ''
    ORG = boto3.client('organizations')
    try:
        masterID = ORG.list_roots()['Roots'][0]['Arn'].rsplit(':')[4]
        return masterID
    except Exception as e:
        logger.error('This stack runs only on the Master of the AWS Organization')
        return False


def create_horizon_stackset(paramList, AdminRoleARN, ExecRole, cList):
    """ Launch CRWD Discover Stackset on the Master Account """
    CFT = boto3.client('cloudformation')
    result = {}
    if len(paramList):
        try:
            result = CFT.create_stack_set(StackSetName=STACKSETNAME,
                                          Description='Roles for CRWD-Horizon',
                                          TemplateURL=STACKSETURL,
                                          Parameters=paramList,
                                          PermissionModel='SERVICE_MANAGED',
                                          Capabilities=cList,
                                          AutoDeployment={
                                              'Enabled': True,
                                              'RetainStacksOnAccountRemoval': False
                                          }, )
            return result
        except ClientError as e:
            if e.response['Error']['Code'] == 'NameAlreadyExistsException':
                logger.info("StackSet already exists")
                result['StackSetName'] = 'CRWD-ROLES-CREATION'
                return result
            else:
                logger.error("Unexpected error: %s" % e)
                result['Status'] = e
                return result


def get_random_alphanum_string(stringLength=15):
    lettersAndDigits = string.ascii_letters + string.digits
    return ''.join((random.choice(lettersAndDigits) for i in range(stringLength)))


def delete_stackset(stacksetName):
    CFT = boto3.client('cloudformation')
    try:
        stackset_result = CFT.describe_stack_set(StackSetName=stacksetName)
        if stackset_result and 'StackSet' in stackset_result:
            stackset_instances = CFT.list_stack_instances(StackSetName=stacksetName)
            while 'NextToken' in stackset_instances:
                stackinstancesnexttoken = stackset_instances['NextToken']
                morestackinstances = CFT.list_stack_instances(NextToken=stackinstancesnexttoken)
                stackset_instances["Summaries"].extend(morestackinstances["Summaries"])
            if len(stackset_instances["Summaries"]) > 0:
                stack_instance_members = [x["Account"] for x in stackset_instances["Summaries"]]
                stack_instance_regions = list(set(x["Region"] for x in stackset_instances["Summaries"]))
                CFT.delete_stack_instances(
                    StackSetName=stacksetName,
                    Accounts=stack_instance_members,
                    Regions=stack_instance_regions,
                    OperationPreferences={'MaxConcurrentCount': 3},
                    RetainStacks=False
                )
            stackset_instances = CFT.list_stack_instances(StackSetName=stacksetName)
            counter = 2
            while len(stackset_instances["Summaries"]) > 0 and counter > 0:
                logger.info("Deleting stackset instance from {}, remaining {}, "
                            "sleeping for 10 sec".format(stacksetName, len(stackset_instances["Summaries"])))
                time.sleep(10)
                counter = counter - 1
                stackset_instances = CFT.list_stack_instances(StackSetName=stacksetName)
            if counter > 0:
                CFT.delete_stack_set(StackSetName=stacksetName)
                logger.info("StackSet {} deleted".format(stacksetName))
            else:
                logger.info("StackSet {} still has stackset instance, skipping".format(stacksetName))
            return True

    except ClientError as e:
        if e.response['Error']['Code'] == 'StackSetNotFoundException':
            logger.info("StackSet {} does not exist".format(stacksetName))
            return True
        else:
            logger.error("Unexpected error: %s" % e)
            return False


def deregister_falcon_horizon_account() -> dict:
    account_id = get_master_id()
    organizational_unit_id = get_org_id()
    auth_token = get_cs_token()
    cspm_api = CSPM.CSPM_Registration(access_token=auth_token)
    # ids=account_id,
    try:
        response = cspm_api.DeleteCSPMAwsAccount(ids=account_id, organization_ids=organizational_unit_id)
        logger.info('Response to register = {}'.format(response))
        if response["status_code"] == 201:
            logger.info('Account Registered')
            return response['body']['resources'][0]
        elif response["status_code"] == 409:
            logger.info('Account already registered - nothing to do')
            return response['resources'][0]
        else:
            error_code = response["status_code"]
            error_msg = response["errors"][0]["message"]
            logger.info('Account Registration Failed - Response {} {}'.format(error_code, error_msg))
            return response['resources'][0]
    except Exception as e:
        logger.info('Got exception {}'.format(e))


def register_falcon_horizon_account() -> dict:
    account_id = get_master_id()
    organizational_unit_id = get_org_id()
    api_message = format_registration_message(account_id, organizational_unit_id)
    auth_token = get_cs_token()
    cpsm_api = CSPM.CSPM_Registration(access_token=auth_token)

    # _result = register_falcon_discover_account(falcon_discover, api_message)
    try:
        response = cpsm_api.CreateCSPMAwsAccount(api_message)
        logger.info('Response to register = {}'.format(response))
        if response["status_code"] == 201:
            logger.info('Account Registered')
            return response['body']['resources'][0]
        elif response["status_code"] == 409:
            logger.info('Account already registered - nothing to do')
            return response['resources'][0]
        else:
            error_code = response["status_code"]
            error_msg = response["errors"][0]["message"]
            logger.info('Account Registration Failed - Response {} {}'.format(error_code, error_msg))
            return response['resources'][0]
    except Exception as e:
        logger.info('Got exception {}'.format(e))


def get_secret(secret_key):
    # Create a Secrets Manager client
    SECRET_REGION = CT_REGION
    session = boto3.session.Session()
    client = session.client(
        service_name='secretsmanager',
        region_name=CT_REGION)
    try:
        get_secret_value_response = client.get_secret_value(
            SecretId=SECRET_NAME)
        logger.debug('Got response code {} calling get_secret with key {}'.format(
            get_secret_value_response['ResponseMetadata']['HTTPStatusCode'], secret_key))
        if 'SecretString' in get_secret_value_response:
            secret_dict = json.loads(get_secret_value_response['SecretString'])
            logger.info('Got secret {} '.format(secret_key))
            return secret_dict[secret_key]
        else:
            return None

    except Exception as e:
        logger.info('Got Error {} retieving secret'.format(e))


def get_cs_token():
    try:
        region_name = CT_REGION
        client_id = get_secret('FalconClientId')
        secret = get_secret('FalconSecret')
        authorized = FalconAuth.OAuth2(
            creds={
                'client_id': client_id,
                'client_secret': secret
            })
    except Exception as e:
        # We can't communicate with the endpoint, return a false token
        logger.info('Got Exception {} generating oauth2 token'.format(e))
        authorized.token = lambda: False
        # Try to retrieve a token from our authentication, returning false on failure
    try:
        token = authorized.token()["body"]["access_token"]
        return token
    except:
        logger.info('Unable to generate authorized token for CS API')
        # Confirm the token was successfully retrieved


def get_org_id():
    ''' Get the organization ID for this account'''
    ORG = boto3.client('organizations')
    orgId = ORG.describe_organization()['Organization']['Id']
    return (orgId)


def format_registration_message(account_id, organization_id):
    data = {
        "resources": [
            {
                "organization_id": organization_id,
                "cloudtrail_region": CT_REGION,
                "account_id": account_id,
            }
        ]
    }
    logger.info('Post Data {}'.format(data))
    return data


def format_deregistration_message(account_id, organization_id):
    data = {
        "resources": [
            {
                "organization_id": organization_id,
                "account_id": account_id,
            }
        ]
    }
    logger.info('Post Data {}'.format(data))
    return data

def create_ioa_stack (paramList):

    CFT = boto3.client('cloudformation')
    cft_result = CFT.create_stack(
        StackName=STACKNAME,
        TemplateURL=CF_NO_CT,
        Parameters=paramList,
        TimeoutInMinutes=5,
        Capabilities=[
            'CAPABILITY_NAMED_IAM',
        ],
        Tags=[
            {
                'Key': 'Vendor',
                'Value': 'CrowdStrike'
            },
        ],
    )
    if cft_result:
        logger.info('Created Stack {}'.format(cft_result.get('StackId')))
        return True
    else:
        return False

def check_stack_complete():
    CFT = boto3.client('cloudformation')


    try:
        # Monitor the status of the stack.   We are looking for 'CURRENT'
        state = 'UNKNOWN'
        timeout = 300  # [seconds]
        timeout_start = time.time()
        while time.time() < timeout_start + timeout:
            time.sleep(10)
            if state == 'CREATE_COMPLETE':
                # Stack has been created.
                return True
            else:
                stack_status = CFT.describe_stacks(StackName=STACKNAME)
                state = stack_status['Stacks'][0]['StackStatus']
        return False
    except Exception as e:
        logger.info('Unable to determine ')
        return False


def lambda_handler(event, context):
    try:

        AccountId = get_master_id()

        #
        # Moved to virtual hosted-style URLs.
        # See https://aws.amazon.com/fr/blogs/aws/amazon-s3-path-deprecation-plan-the-rest-of-the-story/
        # path-style URLs to be depricated
        #

        cList = ['CAPABILITY_IAM', 'CAPABILITY_NAMED_IAM', 'CAPABILITY_AUTO_EXPAND']
        ExecRole = 'AWSControlTowerExecution'
        AdminRoleARN = 'arn:aws:iam::' + AccountId + ':role/service-role/AWSControlTowerStackSetRole'
        logger.info('EVENT Received: {}'.format(event))

        response_data = {}

        if event['RequestType'] in ['Create']:
            logger.info('Event = ' + event['RequestType'])

            register_result = register_falcon_horizon_account()
            logger.info('Account registration result: {}'.format(register_result))
            external_id = register_result['external_id']
            RoleName = register_result['iam_role_arn'].split('/')[-1]
            # Parameters for CRWD-Discover stackset
            horizon_stackset_paramList = []
            keyDict = {}

            keyDict['ParameterKey'] = 'ExternalID'
            keyDict['ParameterValue'] = register_result['external_id']
            horizon_stackset_paramList.append(dict(keyDict))

            keyDict['ParameterKey'] = 'RoleName'
            keyDict['ParameterValue'] = RoleName
            horizon_stackset_paramList.append(dict(keyDict))

            keyDict['ParameterKey'] = 'CSRoleName'
            keyDict['ParameterValue'] = CSRoleName
            horizon_stackset_paramList.append(dict(keyDict))

            keyDict['ParameterKey'] = 'CSAccountNumber'
            keyDict['ParameterValue'] = CSAccountNumber
            horizon_stackset_paramList.append(dict(keyDict))
            logger.debug('CRWD_Horizon Stackset ParamList:{}'.format(horizon_stackset_paramList))
            # Create the Stackset
            CRWD_Horizon_result = create_horizon_stackset(horizon_stackset_paramList, AdminRoleARN,
                                                ExecRole, cList)
            logger.info('CRWD-Horizon Stackset: {}'.format(CRWD_Horizon_result))

            horizon_ioa_stack_paramList = []
            IOAkeyDict = {}

            IOAkeyDict['ParameterKey'] = 'ExternalID'
            IOAkeyDict['ParameterValue'] = register_result['external_id']
            horizon_ioa_stack_paramList.append(dict(IOAkeyDict))

            IOAkeyDict['ParameterKey'] = 'RoleName'
            IOAkeyDict['ParameterValue'] = RoleName
            horizon_ioa_stack_paramList.append(dict(IOAkeyDict))

            IOAkeyDict['ParameterKey'] = 'CSEventBusName'
            IOAkeyDict['ParameterValue'] = register_result['eventbus_name']
            horizon_ioa_stack_paramList.append(dict(IOAkeyDict))

            IOAkeyDict['ParameterKey'] = 'CSBucketName'
            IOAkeyDict['ParameterValue'] = register_result['aws_cloudtrail_bucket_name']
            horizon_ioa_stack_paramList.append(dict(IOAkeyDict))

            IOAkeyDict['ParameterKey'] = 'CSRoleName'
            IOAkeyDict['ParameterValue'] = CSRoleName
            horizon_ioa_stack_paramList.append(dict(IOAkeyDict))

            IOAkeyDict['ParameterKey'] = 'CSAccountNumber'
            IOAkeyDict['ParameterValue'] = CSAccountNumber
            horizon_ioa_stack_paramList.append(dict(IOAkeyDict))

            ioa_stack_id = create_ioa_stack(horizon_ioa_stack_paramList)
            logger.info('Created Stack with id {}'.format(ioa_stack_id))
            stack_complete = check_stack_complete()
            if CRWD_Horizon_result and stack_complete:
                cfnresponse_send(event, context, SUCCESS, CRWD_Horizon_result, "CustomResourcePhysicalID")
                return
            else:
                cfnresponse_send(event, context, FAILED, CRWD_Horizon_result, "CustomResourcePhysicalID")
                return

        elif event['RequestType'] in ['Update']:
            logger.info('Event = ' + event['RequestType'])

            cfnresponse_send(event, context, 'SUCCESS', response_data, "CustomResourcePhysicalID")
            return

        elif event['RequestType'] in ['Delete']:
            logger.info('Event = ' + event['RequestType'])
            deregister_falcon_horizon_account()
            delete_stackset(STACKSETNAME)
            response_data["Status"] = "Success"
            cfnresponse_send(event, context, 'SUCCESS', response_data, "CustomResourcePhysicalID")
            return
        raise Exception
    except Exception as e:
        logger.error(e)
        response_data = {"Status": str(e)}
        cfnresponse_send(event, context, 'FAILED', response_data, "CustomResourcePhysicalID")
        return


if __name__ == '__main__':
    event = {
        "RequestType": "Create",
        "ResponseURL": "http://pre-signed-S3-url-for-response",
        "StackId": "arn:aws:cloudformation:eu-west-1:123456789012:stack/MyStack/guid",
        "RequestId": "unique id for this create request",
        "ResourceType": "Custom::TestResource",
        "LogicalResourceId": "MyTestResource",
        "ResourceProperties": {
            "StackName": "MyStack",
            "List": [
                "1",
                "2",
                "3"
            ]
        }
    }
    context = ()
    lambda_handler(event, context)
