import json
import boto3

def lambda_handler(event, context):
    
    iam = boto3.client('iam')
    ses = boto3.client('ses')
    sns = boto3.client('sns')
    
    user_name = event["detail"]["userIdentity"]["userName"]
    user_policies = iam.list_attached_user_policies(UserName = user_name)
    tags = iam.list_user_tags(UserName = user_name)
    mfa = iam.list_mfa_devices(UserName = user_name)
    
    # Default email variables.
    email_source = 'danie.esquijerosa001@gmail.com'
    email_to = ''
    email_subject = 'MFA Enforcement Tool'
    email_body = ''
    
    # boolean variables to determine whether or not conditions are met
    policy_exists = 0
    mfa_compliant = 0
    email_exists = 0
    
    # Various policy related variables: name, arn, and status attached or not
    policy_name = "Force_MFA"
    policy_arn = 'arn:aws:iam::266945937434:policy/Force_MFA'
    policy_status = 'NOT ATTACHED'
    
    for user_policy in user_policies["AttachedPolicies"]:
        policy_status = 'NOT ATTACHED'
        
        if user_policy["PolicyName"] == policy_name:
            policy_exists = 1
            policy_status = 'POLICY ALREADY EXISTS FOR USER'
            break
        else:
            policy_exists = 0
    
    if policy_exists == 0:
        policy_attachment = iam.attach_user_policy(UserName = user_name, PolicyArn = policy_arn)
        policy_status = "ATTACHED POLICY TO " + user_name   
        
    if len(mfa["MFADevices"]) < 1:
        mfa_compliant = 0
    else:
        mfa_compliant = 1
    
    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=email_source,
                Destination={'ToAddresses': [email_to]},
                Message = {
                    'Subject': {'Data': email_subject},
                    'Body': {'Text': {'Data': email_body}}
                }
            )
            
            #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 user %s is not compliant with MFA Enforcement' % (user_name),
                Subject='MFA Enforcement Tool - Enforcement Update',
            )
            
            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=email_source,
                Destination={'ToAddresses': [email_to]},
                Message = {
                    'Subject': {'Data': email_subject},
                    'Body': {'Text': {'Data': email_body}}
                }
            )
            
            print(ses_respone)
            
            
    
    # TODO implement
    print(policy_status)