import json
import boto3
import os

''' This function collects Ip addresses of running EC2 instances with known prefix. Collected information sends to AWS support and stores into DynamoDB.
- The variables below are set when the lambda function is run.
- This lambda function currently is run manually when a pen test is required on this account.
'''
'''
# email_send_cc = os.environ['email_send_cc']
scan_from_ip_list  = 'lalala'
# time to start testing
start_time = 'lalla'
# time to finish
end_time = 'lalal'
# EC2 instance prefix name. IE 'qa'
instance_prefix = 'zzzzz'
# arn of SES validated email
email_arn = 'arn:aws:ses:us-east-1:497749523224:identity/Andrey.Makarov@entrustdatacard.com'
email_pentest_from = 'Andrey.Makarov@entrustdatacard.com'
submitter_name = 'Andrey.Makarov@entrustdatacard.com'
email_submitter = 'Andrey.Makarov@entrustdatacard.com'
email_additional = 'Andrey.Makarov@entrustdatacard.com'
email_send_from = 'Andrey.Makarov@entrustdatacard.com'
email_send_to = 'Andrey.Makarov@entrustdatacard.com'
email_send_cc = 'Andrey.Makarov@entrustdatacard.com'
'''
# env variables for lambda
# list of IP adresses to scan from
scan_from_ip_list  = os.environ['scan_from_ip_list']
# time to start testing
start_time = os.environ['start_time']
# time to finish
end_time = os.environ['end_time']
# EC2 instance prefix name. IE 'qa'
instance_prefix = os.environ['instance_prefix']
# arn of SES validated email
email_from_arn = os.environ['email_from_arn']
email_pentest_from = os.environ['email_pentest_from']
submitter_name = os.environ['submitter_name']
email_submitter = os.environ['email_submitter']
email_additional = os.environ['email_additional']
email_send_from = os.environ['email_send_from']
email_send_to = os.environ['email_send_to']
email_send_cc = os.environ['email_send_cc']
# aws-security-cust-pen-test@amazon.com


# scans for instances
def run_scan():
    client = boto3.client('ec2')
    response = client.describe_instances(
        DryRun=False,
        InstanceIds=[],
        Filters=[],
    )
    case_list1 = []
    case_list2 = []
    for_email = []
# scans all the instances, then finds all online instances, sorts 'instance_prefix'
    for r in response['Reservations']:
        for i in r[u'Instances']:
            tag_name = i[u'Tags']
            for tag in tag_name:
                if tag[u'Value'].find(instance_prefix) >= 0:
                    info1 = ['\n', 'Instance name: ', i[u'InstanceId'], ' Type:', i[u'InstanceType'], ' Image ID:',
                                 i[u'ImageId'],
                                 '\n', network_interfaces(i['InstanceId']), '\n']
                    for_email.append(info1)
                    case1 = i['InstanceId']
                    case_list1.append(case1)
                    case2 = i['InstanceType']
                    case_list2.append(case2)
                elif tag[u'Key'].find(instance_prefix) >= 0:
                    info1 = ['\n', 'Instance name: ', i[u'InstanceId'], ' Type:', i[u'InstanceType'], ' Image ID:',
                                 i[u'ImageId'],
                                 '\n', network_interfaces(i['InstanceId']), '\n']
                    for_email.append(info1)
                    case1 = i['InstanceId']
                    case_list1.append(case1)
                    case2 = i['InstanceType']
                    case_list2.append(case2)
                break

    for_db = json.dumps({"InstanceId": case_list1, "InstanceType": case_list2}, sort_keys=False,
                        indent=4, separators=(',', ': '))
    # sends email
    message = ''
    client = boto3.client("ses")
    for i in for_email:
        # ''.join(list(i[7]))
        ii = ''.join(list(i))
        message += ii
    # message = repr(for_email)
    print(message)
    session = boto3.session.Session()
    region = session.region_name
    token = boto3.client("sts")
    account_id = token.get_caller_identity()["Account"]
    iplist = 'Source IPs: ' + scan_from_ip_list
    info = 'Email Address: {0},\nAWS AccountId: {1},\nSubmitterName: {2},\nCompanyName: EntrustDatacard Corp,' \
           '\nEmailAddress: {3},\n AdditionalEmail1: {4},\n IPs to be scanned: \n'.format(email_pentest_from, account_id, submitter_name, email_submitter, email_additional)
    rowMess = 'From: {0}\nTo: {1}\nCC: {2}\nSubject: Your AWS ' \
              'Penetration Testing Inquiry\nMIME-Version: 1.0\n\n{3}\n{4}'.format(email_send_from, email_send_to, email_send_cc, info, message) + '\n\nBandwidth\t\t\t1\nRegion\t\t    ' \
                '\t' + region + '\nTimezone\t\t\tgmt--3\nStartDate\t\t\t' + start_time +'\nEndDate\t\t\t'+ end_time +'\nComments\t\t' \
                '\t\nTermsAgreement\t\t\ti-agree\nPolicy\t\t\ti-agree\nAWS Account ID\t\t\t\nIAM ' \
                                'Account\t\t\tNo\nEmail ' \
                'Address\t\t\t'+ email_pentest_from +'\nName\t' \
                '\t\tacs_production \n' + iplist

    response = client.send_raw_email(
        Destinations=[
        ],
        FromArn='',
        RawMessage={
            'Data': rowMess,
        },
        SourceArn='' + email_from_arn,
    )

    print(response)

    return for_db


# updates DynamoDB table
def load_data(for_db):
    import json
    import decimal
    dynamodb = boto3.resource('dynamodb')
    import datetime
    my_time = str(datetime.datetime.now())
    print(my_time)
    table = dynamodb.Table('Instances')

    opened_json = json.loads(for_db)
    for i in opened_json['InstanceId']:
        InstanceId = i
        InstanceId = str(InstanceId)
        for k in opened_json['InstanceType']:
            InstanceType = k
            InstanceType = str(InstanceType)
        print("Adding InstanceId:", InstanceId)

        status = table.put_item(
            Item={
                'TimeStamp': my_time,
                'InstanceId': InstanceId,
            }
        )
        print(status)

#function to scan online instance
def network_interfaces(instance):
    s3 = boto3.resource('s3')
    client = boto3.client('ec2')
    ni_responce = str
    response = client.describe_instances(
        DryRun=False,
        InstanceIds=[
            instance,
        ],
        Filters=[],
    )
    for r in response['Reservations']:
        for k in r['Instances']:
            if k[u'State'][u'Name'] == 'running':
                for i in k['NetworkInterfaces']:
                    try:
                        ni_responce = ''.join(['Private IP:', i['PrivateIpAddress'],' Instance is Running. Public IP:', i['Association']['PublicIp']])
                    except:
                        ni_responce = 'no public IP'
            else:
                ni_responce = 'Instance is not running'
    return ni_responce


# function to create DynamoDB table
def create_table():
    dynamodb = boto3.resource('dynamodb')
    try:
        response = dynamodb.create_table(
            AttributeDefinitions=[
                {
                    'AttributeName': 'TimeStamp',
                    'AttributeType': 'S',
                },
                {
                    'AttributeName': 'InstanceId',
                    'AttributeType': 'S',
                },
            ],
            KeySchema=[
                {
                    'AttributeName': 'TimeStamp',
                    'KeyType': 'HASH',
                },
                {
                    'AttributeName': 'InstanceId',
                    'KeyType': 'RANGE',
                },
            ],
            ProvisionedThroughput={
                'ReadCapacityUnits': 5,
                'WriteCapacityUnits': 5,
            },
            TableName='Instances',
        )
        print("Table status:", response.table_status)
    except:
        print ('Table exists')


# standart lambda handler from AWS
def lambda_handler(event, context):
    for_db = run_scan()
    create_table()
    load_data(for_db)
