import json
import boto3

ec2_client = boto3.client('ec2')
s3_client = boto3.client('s3')
rds_client = boto3.client('rds')

tag_to_check = 'Owner'

def check_for_tag(my_list):
    print('Looking for tag [' + tag_to_check + '] in list ' + json.dumps(my_list))
    for i in my_list:
        if i['Key'] == tag_to_check:
            return True
    return False

def tag_bucket(bucket, user):
    s3_client.put_bucket_tagging(
        Bucket=bucket,
        Tagging={
            'TagSet': [
                {
                    'Key': tag_to_check,
                    'Value': user
                },
            ]
        }
    )

def tag_instance(instance, user):
    ec2_client.create_tags(
        Resources=[instance],
        Tags=[
            {
                'Key': tag_to_check,
                'Value': user
            }
        ]
    )

def tag_rds_instance(rds_instance, user):
    rds_client.add_tags_to_resource(
        ResourceName=rds_instance,
        Tags=[
            {
                'Key': tag_to_check,
                'Value': user
            }
        ]
    )

def lambda_handler(event, context):
    print(json.dumps(event))

    user = ''
    try:
        # IAM user
        user = event['detail']['userIdentity']['userName']
    except:
        user = event['detail']['userIdentity']['arn'].rsplit('/', 1)[-1]

    if event['source'] == "aws.s3":
        bucket = event['detail']['requestParameters']['bucketName']

        # check for existing tags
        try:
            bucket_tags = json.loads(s3_client.get_bucket_tagging(Bucket=bucket))
            if check_for_tag(bucket_tags['TagSet']):
                print('Bucket [' + bucket + '] is already tagged with [' + tag_to_check + ']')
                return
            else:
                print('Tagging Bucket [' + bucket + '] with ' + tag_to_check + '=' + user)
                tag_bucket(bucket, user)
                return
        except:
            # if an exception is raised, the bucket is not tagged
            print('No tags found on bucket [' + bucket + ']')
            print('Tagging Bucket [' + bucket + '] with ' + tag_to_check + '=' + user)
            tag_bucket(bucket, user)
            return

    elif event['source'] == "aws.ec2":
        instances = event['detail']['responseElements']['instancesSet']['items']
        instances_list = []
        for j in instances:
            instances_list.append(j['instanceId'])
        print(str(instances_list))
        response = ec2_client.describe_instances(InstanceIds=instances_list)
        for i in instances_list:
            print(i)
            try:
                if check_for_tag(i['Tags']):
                    print('Instance [' + i + '] is already tagged with [' + tag_to_check + ']')
                    return
                else:
                    print('Tagging Instance [' + i + '] with ' + tag_to_check + '=' + user)
                    tag_instance(i, user)
                    return                    
            except:
                # if an exception is raised, the instance is not tagged
                print('No tags found on instance [' + i + ']')
                print('Tagging Instance [' + i + '] with ' + tag_to_check + '=' + user)
                tag_instance(i, user)
                return
            
    elif event['source'] == "aws.rds":
        rds_instance = event['detail']['responseElements']['dBInstanceArn']
        rds_tags = rds_client.list_tags_for_resource(ResourceName=rds_instance)
        try:
            if check_for_tag(rds_tags['TagList']):
                print('RDS Instance [' + rds_instance + '] is already tagged with [' + tag_to_check + ']')
                return
            else:
                print('Tagging RDS Instance [' + rds_instance + '] with ' + tag_to_check + '=' + user)
                tag_rds_instance(rds_instance, user)
                return                    
        except:
            # if an exception is raised, the instance is not tagged
            print('No tags found on RDS instance [' + rds_instance + ']')
            print('Tagging RDS Instance [' + rds_instance + '] with ' + tag_to_check + '=' + user)
            tag_rds_instance(rds_instance, user)
            return


    return