aws_project_initiation.py

import boto3
import json
import string
import random
import warnings

from botocore.config import Config
from ds_helpers.aws import get_secrets_manager_secret
from ds_helpers.db import connect_to_mysql

warnings.filterwarnings('ignore')


def generate_random_password(length=21):
    characters = list(string.ascii_letters + string.digits + "!#$")
    random.shuffle(characters)
    password = []
    for i in range(length):
        password.append(random.choice(characters))
    random.shuffle(password)
    return "".join(password)


def main(project_name, s3_bucket_names, secrets, prod_mysql_secret, stage_mysql_secret, ecr_policy_arn, s3_policy_arn,
         secrets_manager_policy_arn, region='us-east-1'):
    boto3_config = Config(
        region_name=region,
    )
    arn_list = []

    codecommit_client = boto3.client('codecommit', config=boto3_config)
    codecommit_client.create_repository(
        repositoryName=project_name,
    )

    ecr_client = boto3.client('ecr', config=boto3_config)
    ecr_response = ecr_client.create_repository(
        repositoryName=project_name,
        imageTagMutability='MUTABLE',
        imageScanningConfiguration={
            'scanOnPush': True
        },
    )
    arn_list.append(('ecr', project_name, ecr_response.get('repository').get('repositoryArn')))

    s3_client = boto3.client('s3', config=boto3_config)
    for bucket in s3_bucket_names:
        temp_s3_response = s3_client.create_bucket(
            ACL='private',
            Bucket=bucket,
        )
        arn_list.append(('s3', bucket, f'arn:aws:s3:::{bucket}'))

    iam_client = boto3.client('iam', config=boto3_config)
    iam_client.create_user(
        UserName=project_name
    )
    iam_response = iam_client.create_access_key(
        UserName=project_name
    )
    key_id = iam_response['AccessKey'].get('AccessKeyId')
    secret_key = iam_response['AccessKey'].get('SecretAccessKey')

    secrets_manager_client = boto3.client("secretsmanager", config=boto3_config)
    secrets_manager_response = secrets_manager_client.create_secret(Name=f'aws_keys_{project_name}',
                                                                    SecretString=json.dumps({
                                                                        "access_key": key_id,
                                                                        "secret_key": secret_key,
                                                                    }))
    arn_list.append(('secrets_manager', f'aws_keys_{project_name}', secrets_manager_response.get('ARN')))
    for key, value in secrets.items():
        temp_secrets_manager_response = secrets_manager_client.create_secret(Name=key, SecretString=value)
        arn_list.append(('secrets_manager', key, temp_secrets_manager_response.get('ARN')))

    def create_mysql_structure(db_secret, environment):
        mysql_project_name = project_name.replace('-', '_')
        secret_dict = get_secrets_manager_secret(db_secret)
        db_conn = connect_to_mysql(secret_dict, ssl_path='../../../../data/rds-ca-2019-root.pem')
        db_conn.execute(f'''create database if not exists {mysql_project_name};''')
        db_conn.execute(f'''use {mysql_project_name};''')
        db_conn.execute(f'''
        create procedure CreateModelingTables()

        begin

        create table if not exists {mysql_project_name}.model_score (
            id int not null primary key auto_increment,
            meta__inserted_at timestamp default current_timestamp,
            model_uid text,
            holdout_score float,
            scoring_metric text,
            holdout_type enum('test', 'cross_validation')
            );

        create table if not exists {mysql_project_name}.feature_score (
            id int not null primary key auto_increment,
            meta__inserted_at timestamp default current_timestamp,
            model_uid text,
            feature text,
            importance_score float,
            importance_metric text
            );

        create table if not exists {mysql_project_name}.model_metadata (
            id int not null primary key auto_increment,
            training_timestamp timestamp,
            model_uid text
        );

        create table if not exists {mysql_project_name}.model_logs (
            id int not null auto_increment primary key,
            meta__inserted_at timestamp not null default current_timestamp,
            uid text,
            logging_timestamp timestamp,
            input_output_payloads json
            );

        create table if not exists {mysql_project_name}.config (
            id int not null auto_increment primary key,
            meta__inserted_at timestamp not null default current_timestamp,
            config_key text,
            config_value text
            );    

        end
        ''')
        db_conn.execute('''call CreateModelingTables();''')
        password = generate_random_password()
        db_conn.execute(f'''create user if not exists '{mysql_project_name}'@'%%' identified by '{password}';''')
        db_conn.execute(
            f'''grant insert, select, update, delete, create on {mysql_project_name}.* to '{mysql_project_name}'@'%%';''')
        db_conn.execute('''flush privileges;''')
        sm_response = secrets_manager_client.create_secret(Name=f'{environment}-mysql-{mysql_project_name}',
                                                           SecretString=json.dumps({
                                                               "host": secret_dict.get('host'),
                                                               "database": mysql_project_name,
                                                               "user": mysql_project_name,
                                                               "password": password
                                                           }))
        return sm_response

    sm_stage_response = create_mysql_structure(stage_mysql_secret, 'stage')
    arn_list.append(('secrets_manager', f'stage_mysql_{project_name}', sm_stage_response.get('ARN')))
    sm_prod_response = create_mysql_structure(prod_mysql_secret, 'prod')
    arn_list.append(('secrets_manager', f'prod_mysql_{project_name}', sm_prod_response.get('ARN')))

    iam_client.create_group(
        GroupName=project_name
    )
    iam_client.add_user_to_group(
        GroupName=project_name,
        UserName=project_name
    )

    for arn in arn_list:
        if arn[0] == 'ecr':
            policy = json.dumps({
                "Version": "2012-10-17",
                "Statement": [
                    {
                        "Effect": "Allow",
                        "Action": "ecr:GetAuthorizationToken",
                        "Resource": "*"
                    },
                    {
                        "Effect": "Allow",
                        "Action": "ecr:*",
                        "Resource": [
                            f"{arn[2]}"
                        ]
                    }
                ]
            })
            iam_client.create_policy(
                PolicyName=f'{arn[1]}-ecr',
                PolicyDocument=policy,
            )
            main_policy = iam_client.get_policy(
                PolicyArn=ecr_policy_arn
            )
            main_policy_version = iam_client.get_policy_version(
                PolicyArn=ecr_policy_arn,
                VersionId=main_policy['Policy']['DefaultVersionId']
            )
            main_policy_version['PolicyVersion']['Document']['Statement'][1]['Resource'] = \
                main_policy_version['PolicyVersion']['Document']['Statement'][1]['Resource'] + [arn[2]]
            main_policy_version = main_policy_version.get('PolicyVersion').get('Document')
            iam_client.create_policy_version(
                PolicyArn=ecr_policy_arn,
                PolicyDocument=json.dumps(main_policy_version, indent=4, sort_keys=True, default=str),
                SetAsDefault=True
            )
        elif arn[0] == 's3':
            policy = json.dumps({
                "Version": "2012-10-17",
                "Statement": [
                    {
                        "Effect": "Allow",
                        "Action": "s3:*Object",
                        "Resource": [
                            f"{arn[2]}/*"
                        ]
                    }
                ]
            })
            response = iam_client.create_policy(
                PolicyName=f'{arn[1]}-s3',
                PolicyDocument=policy,
            )
            iam_client.attach_group_policy(
                GroupName=project_name,
                PolicyArn=response.get('Policy').get('Arn')
            )
            main_policy = iam_client.get_policy(
                PolicyArn=s3_policy_arn
            )
            main_policy_version = iam_client.get_policy_version(
                PolicyArn=s3_policy_arn,
                VersionId=main_policy['Policy']['DefaultVersionId']
            )
            main_policy_version['PolicyVersion']['Document']['Statement'][0]['Resource'] = \
                main_policy_version['PolicyVersion']['Document']['Statement'][0]['Resource'] + [arn[2] + '/*']
            main_policy_version = main_policy_version.get('PolicyVersion').get('Document')
            iam_client.create_policy_version(
                PolicyArn=s3_policy_arn,
                PolicyDocument=json.dumps(main_policy_version, indent=4, sort_keys=True, default=str),
                SetAsDefault=True
            )
        elif arn[0] == 'secrets_manager':
            policy = json.dumps({
                "Version": "2012-10-17",
                "Statement": [
                    {
                        "Effect": "Allow",
                        "Action": "secretsmanager:GetSecretValue",
                        "Resource": [
                            f"{arn[2]}"
                        ]
                    }
                ]
            })
            response = iam_client.create_policy(
                PolicyName=f'{arn[1]}-secrets-manager',
                PolicyDocument=policy,
            )
            iam_client.attach_group_policy(
                GroupName=project_name,
                PolicyArn=response.get('Policy').get('Arn')
            )
            if arn[1] != f'aws_keys_{project_name}':
                main_policy = iam_client.get_policy(
                    PolicyArn=secrets_manager_policy_arn
                )
                main_policy_version = iam_client.get_policy_version(
                    PolicyArn=secrets_manager_policy_arn,
                    VersionId=main_policy['Policy']['DefaultVersionId']
                )
                main_policy_version['PolicyVersion']['Document']['Statement'][0]['Resource'] = \
                    main_policy_version['PolicyVersion']['Document']['Statement'][0]['Resource'] + [arn[2]]
                main_policy_version = main_policy_version.get('PolicyVersion').get('Document')
                iam_client.create_policy_version(
                    PolicyArn=secrets_manager_policy_arn,
                    PolicyDocument=json.dumps(main_policy_version),
                    SetAsDefault=True
                )
        else:
            raise Exception('can only create policies for ecr, s3, or secrets manager at this time')