import json
import boto3

def lambda_handler(event, context):
    iam = boto3.client('iam')
    ses = boto3.client('ses')
    sns = boto3.client('sns')
    users = iam.list_users()
    user_list = []
    mfa_compliant = 0
    
    # Default email variables.
    email_to = ''
    email_subject = 'MFA Enforcement Tool'
    email_body = ''
    email_exists = 0
    
    # Loops through all users in IAM
    for user in users["Users"]:
        username = user["UserName"]
        tags = iam.list_user_tags(UserName = username)
        mfa = iam.list_mfa_devices(UserName = username)
        
        # Checks to see if specific user has any MFA devices, and adds user to user_list if not compliant.
        if len(mfa["MFADevices"]) < 1:
            mfa_compliant = 0
            user_list.append({'username': user["UserName"], 'userId': user["UserId"]})
        else:
            mfa_compliant = 1
        
        # Checks to see if user has tags, in this case, their email.
        if len(tags["Tags"]) > 0:
            for tag in tags["Tags"]:
                if tag["Key"] == "email":
                    email_to = tag["Value"]
                    email_exists = 1
                else:
                    email_exists = 0

        # If not compliant, send an email with steps to take in order to become compliant.
        if email_exists == 1:
            if mfa_compliant == 0:
                
                email_body = '''This is an automated message.\n\nYou are not MFA compliant. Follow these steps in order to become MFA compliant.
                \n1. Go to your IAM dashboard.
                \n2. Click on your username.
                \n3. Click on "Security Credentials"
                \n4. Click on "Manage MFA Device"
                \n5. Follow the instructions to set up your MFA device.
                '''
                
                ses_respone = ses.send_email(
                    Source='danie.esquijerosa001@gmail.com',
                    Destination={'ToAddresses': [email_to]},
                    Message = {
                        'Subject': {'Data': email_subject},
                        'Body': {'Text': {'Data': email_body}}
                    }
                )
                print(ses_respone)
            # If compliant, thanks them for being compliant.    
            else:
                email_body = '''This is an automated message.\n\nThank you for being MFA compliant. No further intervention is required.'''
                ses_respone = ses.send_email(
                    Source='danie.esquijerosa001@gmail.com',
                    Destination={'ToAddresses': [email_to]},
                    Message = {
                        'Subject': {'Data': email_subject},
                        'Body': {'Text': {'Data': email_body}}
                    }
                )
                print(ses_respone)
        
        # Sets user dependent variables back to default.
        email_to = ''
        email_exists = 0
    
    #Pusblishes results to SNS topic sent directly to administrators.
    sns_response = sns.publish(
        TopicArn='arn:aws:sns:us-east-1:266945937434:MFA_Enforcement_Alerts',    
        Message='The following users are not compliant with MFA Enforcement: %s' % (user_list),
        Subject='MFA Enforcement Tool - Enforcement Update',
    )
    
    print(sns_response)