import json
import boto3
import os

cloudtrail_client = boto3.client('cloudtrail')
ssm_client = boto3.client('ssm')
sqs_client = boto3.client('sqs')

def lambda_handler(event, context):
    
    print(json.dumps(event))
    # get environment variables
    jira_username = os.environ['JIRA_USERNAME']
    jira_password_parameter_name = os.environ['JIRA_PASSWORD_PARAMETER_NAME']
    jira_project_key = os.environ['JIRA_PROJECT_KEY']
    jira_url = os.environ['JIRA_URL']
    jira_issue_type_name = os.environ['JIRA_ISSUE_TYPE_NAME']
    
    delay_seconds = 300

    print('Querying CloudTrail for RunInstances events')
    cloudtrail_events = cloudtrail_client.lookup_events(
        LookupAttributes=[
            {
                'AttributeKey': 'EventName',
                'AttributeValue': 'RunInstances'
            }
        ]
    )

    for record in event['Records']:
        receiptHandle = record['receiptHandle']
        body = json.loads(record['body'])
        queue_arn = record['eventSourceARN']
        issue_summary = 'EBS Volume ' + body['volume'] + ' is unencrypted'
        issue_description = ''

        print(body)
        # count how many events matched this instance
        count = 0

        for e in cloudtrail_events['Events']:
            for r in e['Resources']:
                if r['ResourceName'] == body['instance']:
                    # found a RunInstances event that matched this instance
                    count += 1
                    # instead of just printing 'AutoScaling', print the IAM Role that initiated the AutoScaling event
                    if e['Username'] == 'AutoScaling':
                        print('IAM Role ' + json.loads(e['CloudTrailEvent'])['userIdentity']['arn'] + ' created instance ' + body['instance'] + ' in account ' + body['account'] + ' with unencrypted volume ' + body['volume'])
                        issue_description = 'IAM Role ' + json.loads(e['CloudTrailEvent'])['userIdentity']['arn'] + ' created instance ' + body['instance'] + ' in account ' + body['account'] + ' with unencrypted volume ' + body['volume']
                    else:
                        print('User ' + e['Username'] + ' created instance ' + body['instance'] + ' in account ' + body['account'] + ' with unencrypted volume ' + body['volume'])
                        issue_description = 'User ' + e['Username'] + ' created instance ' + body['instance'] + ' in account ' + body['account'] + ' with unencrypted volume ' + body['volume']
        
        # parse the queue name from the arn
        queue_name = queue_arn.rsplit(':', 1)[-1]
        # get the the queue url
        queue_url = sqs_client.get_queue_url(QueueName=queue_name)['QueueUrl']
        # if no RunInstances events matched this instance
        
        if count == 0:
             print('No events found in CloudTrail for creating EC2 instance ' + body['instance'] + ' with unencrypted volume ' + body['volume'] + ', sending back to queue')
             # delete original message from queue, and send a new message to queue to try again later
             sqs_client.delete_message(QueueUrl=queue_url, ReceiptHandle=receiptHandle)
             sqs_client.send_message(QueueUrl=queue_url, DelaySeconds=delay_seconds, MessageBody=json.dumps({"instance": body['instance'], "volume":body['volume'], "account":body['account']}))
        else:
            issue_description += '\\n\\nMore information about encrypting EBS Volumes: https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/EBSEncryption.html'

            response = ssm_client.start_automation_execution(
                DocumentName='AWS-CreateJiraIssue',
                Parameters={
                    'JiraUsername': [jira_username],
                    'SSMParameterName': [jira_password_parameter_name],
                    'JiraURL': [jira_url],
                    'ProjectKey': [jira_project_key],
                    'IssueTypeName': [jira_issue_type_name],
                    'IssueSummary': [issue_summary],
                    'IssueDescription': [issue_description]
                }
            )
            
            print('Sent Automation Execution ID ' + response['AutomationExecutionId'])
            
            # delete the message from the queue
            sqs_client.delete_message(QueueUrl=queue_url, ReceiptHandle=receiptHandle)

    return