import boto3
from botocore.exceptions import ClientError

class Auditor:

    def __init__(self,accounts,regions):
        self.accounts = accounts
        self.regions  = regions
        self.role     = 'SecurityRole'
        self.data     = {}

    # Populate Instance IP Addresses in our Object data struct.
    def popinstanceips(self,dic,client,account,region):
        print("Getting Data Account= " + account + "Region=" + region)
        instances = client.describe_instances()
        instances = instances['Reservations']
        result = [x['Instances'][0] for x in instances if 'PublicIpAddress' in x['Instances'][0]]
        for item in result:
            dic[item['InstanceId']] = {
                'ip':item['PublicIpAddress'],
                'technology':'ec2',
                'account':account,
                'region': region
            }

    def ec2_gather(self,account,region,tokens):
        client = boto3.client('ec2',
                region_name=region,
                aws_access_key_id=tokens[0],
                aws_secret_access_key=tokens[1],
                aws_session_token=tokens[2]
                )
        self.popinstanceips(self.data,client,account,region)

    def assume_role(self,account,region):
        boto_sts = boto3.client('sts')
        arn ='arn:aws:iam::' + account + ':role/' + self.role
        try:
            sts_response = boto_sts.assume_role(RoleArn=arn, RoleSessionName='securitymonkey')
            newsession_id = sts_response["Credentials"]["AccessKeyId"]
            newsession_key = sts_response["Credentials"]["SecretAccessKey"]
            newsession_token = sts_response["Credentials"]["SessionToken"]
            return newsession_id,newsession_key,newsession_token
        except ClientError as e:
            print("Account= " + account + " Could not be assumed by Tool! Error: " + str(e) )
            return None


    def call_sts(self,data,account):
        # enumerate regions.
        for region in self.regions:
            tokens = self.assume_role(account,region)
            if tokens != None:
                self.ec2_gather(account,region,tokens)

    def auditor(self):
        for account in self.accounts:
            self.call_sts(self.data,account)
        print("## Instance IP Report ##")
        print(self.data)

regions = ['us-west-2','us-west-1','us-east-1','us-east-2']
accounts = ['accountnumber1','accountnumber2']

auditor = Auditor(accounts,regions)
auditor.auditor()