####################################################################
##  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
##  SPDX-License-Identifier: MIT-0
####################################################################
##  iam-cred-report-gen-report
##
##      input fields
##          accountId   : required, the AWS account id to process
##          LoopAgain   : default "yes"
##                          Controls the looping in the Step function
##          LoopCount   : default 0
##                          Keeps track of the number of times we have looped
##          funcState   : default "not complete"
##                          A state of "complete" indicates success
##                          A state of "error" indicates an error
##
##      output fields
##          All the input fields
##          reportFileName :    The name of the s3 file created
##                              under the date folder in the bucket
##

import json
import os
import boto3
import time
import logging
import sys
import botocore
import tempfile

from datetime import date


logger = logging.getLogger(__name__)
logger.setLevel(level=logging.INFO)  # use level DEBUG for more details

## Define current date
today = str(date.today())

###############################################################################
##  Get ENV Vars

## Read environment Variables from Lambda Function
BUCKET_ARN = os.getenv('BUCKET_ARN')
if not BUCKET_ARN:
    raise Exception ("Missing BUCKET_ARN value")
logger.debug ("BUCKET_ARN: %s" % (BUCKET_ARN))

ASSUME_ROLE_NAME = os.getenv('ASSUME_ROLE_NAME')
if not ASSUME_ROLE_NAME:
    raise Exception ("Missing ASSUME_ROLE_NAME value")
logger.debug ("ASSUME_ROLE_NAME: %s" % (ASSUME_ROLE_NAME))

MAX_LOOP = os.getenv ('MAX_LOOP')
if not MAX_LOOP:
    MAX_LOOP = '15'
if not MAX_LOOP.isnumeric():
    raise Exception ("MAX_LOOP is not numeric: %s" % (MAX_LOOP))


###############################################################################
## Setup bucket for storing reports and get stsClient
s3 = boto3.resource('s3')
bucketName = BUCKET_ARN.split(":", -1)[-1]
bucketConnection = s3.Bucket(bucketName)
    
stsClient = boto3.client('sts')

###############################################################################
## Function to handle the temp file and S3 bucket upload
def write_to_temp_and_upload(filename,content):
    tmpdir = tempfile.mkdtemp()

    full_path = os.path.join(tmpdir, filename)
    #print(full_path)
    try:
        with open(full_path, "w") as tmp:
            tmp.write(content)
    except IOError as e:
        print('IOError: Unable to write IAM Credential Report to Temp File.')
    finally:
        s3.Object(bucketName, today+"/"+filename).put(Body=open(full_path, 'rb'))
        ## to avoid the lambda tmp space from filling, remove the file we created
        os.remove(full_path)
        ## this is used to remove the empty directory. 
        ## OSError will be raised if the specified path is not an empty directory.
        os.rmdir(tmpdir)
        



###############################################################################
## main

def lambda_handler(event, context):

    ## Log what we were passed
    logger.debug ("Input event: %s" % (json.dumps (event)))
    logger.debug ("MAX_LOOP: %s" % (MAX_LOOP))

    funcStatus = {}

    ## copy event to funcStatus
    funcStatus = event

    ## this is fatal
    if 'accountId' not in funcStatus:
        logger.error ("Error, we don't have an accountId")
        raise Exception ("no accountId was passed to function")

    ## check if we are in a loop, if not, initialize it
    if 'LoopCount' not in funcStatus:
        funcStatus['LoopCount'] = 0
    else:
        ## increment the counter
        funcStatus['LoopCount'] += 1
        
    if 'LoopAgain' not in funcStatus:
        funcStatus['LoopAgain'] = "yes"
        
    if 'funcState' not in funcStatus:
        funcStatus['funcState'] = "not complete"
        
    ## Define current date
    today = str(date.today())

    ## get the current account id and only assume role if what we are processing
    ## is different
    CurrentAccount = stsClient.get_caller_identity().get('Account')
    
    ## process the record
    logger.info("Processing AccountId: %s" % (funcStatus['accountId']))
    
    ## We need to switch to the correct account
    try:
        ## only switch roles if necessary
        if funcStatus['accountId'] != CurrentAccount:
            assumedRoleObject=stsClient.assume_role(
                RoleArn=f"arn:aws:iam::{funcStatus['accountId']}:role/{ASSUME_ROLE_NAME}",
                RoleSessionName="IAMCredentialReport")
            credentials=assumedRoleObject['Credentials']
    
            iamClient=boto3.client('iam',
                aws_access_key_id=credentials['AccessKeyId'],
                aws_secret_access_key=credentials['SecretAccessKey'],
                aws_session_token=credentials['SessionToken'])
        else:
            iamClient=boto3.client('iam')
            
    except botocore.exceptions.ClientError as error:
        logger.error(error)
        raise error

    ## Generate report / check current state of report generation
    try:
        # Generate Credential Report
        reportcomplete = False
        gencredentialreport = iamClient.generate_credential_report()
        logger.info('IAM credential report current state: %s' % (gencredentialreport['State']))
        
        ## set true / false
        reportcomplete = gencredentialreport['State'] == 'COMPLETE'
        #reportcomplete = False  ## For testing

        if not reportcomplete:
            ###########################################################################        
            ## Error out if we have looped too many times
            ## The Step function will handle the error
            if funcStatus['LoopCount'] > int (MAX_LOOP):
                    logger.error ("Error: Too many iterations. Exceeded %s loops." % (MAX_LOOP))
                    funcStatus['ErrorMessage'] = f"Too many iterations. Exceeded MAX_LOOP of {MAX_LOOP}"
                    funcStatus['LoopAgain'] = "no"
                    funcStatus['funcState'] = "error"
                    return (funcStatus)
                    
            ## otherwise, loop
            else:
                logger.info ("Waiting for report to complete... Exiting to sleep in Step function.")
                funcStatus['LoopAgain'] = "yes"
                return (funcStatus)            


        ## report should be complete at this point
        logger.info('IAM credential report successfully generated for account Id: %s' % (funcStatus['accountId']))
        credentialReport = iamClient.get_credential_report()
        decodedCredentialReport = credentialReport['Content'].decode("utf-8")
        ## Save credential Report into CSV file
        funcStatus['reportFileName'] = f"credentialReport_{funcStatus['accountId']}.csv"
        
        # Write report to a temp file and upload to S3
        write_to_temp_and_upload(funcStatus['reportFileName'],decodedCredentialReport)
     
        ## everything should be done by this point
        funcStatus['LoopAgain'] = "no"
        funcStatus['funcState'] = "complete"        

        return (funcStatus)


    except iamClient.exceptions.LimitExceededException as error:
        logger.info ("Waiting for limit... Exiting to sleep in Step function.")
        funcStatus['LoopAgain'] = "yes"
        return (funcStatus) 

    except botocore.exceptions.ClientError as error:
        logger.error(error)
        raise error

    raise Exception ('Error, we should never be here, unknown code path.')