import pulumi
import pulumi_aws as aws
import conf as config
import json


def create_execution_role():
    execution_role = aws.iam.Role(f"{config.project_name}-execution-role",
        assume_role_policy=json.dumps({
            "Statement": [{
                "Action": "sts:AssumeRole",
                "Effect": "Allow",
                "Principal": {
                    "Service": "ecs-tasks.amazonaws.com",
                },
            }],
            "Version": "2012-10-17",
        }),
        inline_policies=[aws.iam.RoleInlinePolicyArgs(
            name=f"{config.project_name}-{config.stack_name}-service-secrets-policy",
            policy=json.dumps({
                "Version": "2012-10-17",
                "Statement": [
                    {
                        "Effect": "Allow",
                        "Action": [
                            "ecr:GetAuthorizationToken",
                            "ecr:BatchCheckLayerAvailability",
                            "ecr:GetDownloadUrlForLayer",
                            "ecr:GetRepositoryPolicy",
                            "ecr:DescribeRepositories",
                            "ecr:ListImages",
                            "ecr:DescribeImages",
                            "ecr:BatchGetImage",
                            "ecr:GetLifecyclePolicy",
                            "ecr:GetLifecyclePolicyPreview",
                            "ecr:ListTagsForResource",
                            "ecr:DescribeImageScanFindings"
                        ],
                        "Resource": "*"
                    }
                ],
            }),
        )],
        managed_policy_arns=["arn:aws:iam::aws:policy/service-role/AmazonECSTaskExecutionRolePolicy"])
    return execution_role

def create_execution_role_with_keys(ssm_parameter, key):
    execution_role = aws.iam.Role(f"{config.project_name}-execution-role",
        assume_role_policy=json.dumps({
            "Statement": [{
                "Action": "sts:AssumeRole",
                "Effect": "Allow",
                "Principal": {
                    "Service": "ecs-tasks.amazonaws.com",
                },
            }],
            "Version": "2012-10-17",
        }),
        inline_policies=[aws.iam.RoleInlinePolicyArgs(
            name=f"{config.project_name}-{config.stack_name}-service-secrets-policy",
            policy=pulumi.Output.all(ssm_parameter.arn, key.arn).apply(lambda args: json.dumps({
                "Version": "2012-10-17",
                "Statement": [
                    {
                        "Action": ["ssm:GetParameters"],
                        "Condition": {
                            "StringEquals": {
                                "ssm:ResourceTag/pulumi-application": config.project_name,
                                "ssm:ResourceTag/pulumi-environment": config.stack_name,
                            },
                        },
                        "Effect": "Allow",
                        "Resource": [args[0]],
                    },
                    {
                        "Action": ["kms:Decrypt"],
                        "Condition": {
                            "StringEquals": {
                                "aws:ResourceTag/pulumi-application": config.project_name,
                                "aws:ResourceTag/pulumi-environment": config.stack_name,
                            },
                        },
                        "Effect": "Allow",
                        "Resource": [args[1]],
                        "Sid": "DecryptTaggedKMSKey",
                    },
                    {
                        "Effect": "Allow",
                        "Action": [
                            "ecr:GetAuthorizationToken",
                            "ecr:BatchCheckLayerAvailability",
                            "ecr:GetDownloadUrlForLayer",
                            "ecr:GetRepositoryPolicy",
                            "ecr:DescribeRepositories",
                            "ecr:ListImages",
                            "ecr:DescribeImages",
                            "ecr:BatchGetImage",
                            "ecr:GetLifecyclePolicy",
                            "ecr:GetLifecyclePolicyPreview",
                            "ecr:ListTagsForResource",
                            "ecr:DescribeImageScanFindings"
                        ],
                        "Resource": "*"
                    }
                ],
            })),
        )],
        managed_policy_arns=["arn:aws:iam::aws:policy/service-role/AmazonECSTaskExecutionRolePolicy"])
    return execution_role


def create_task_role():
    task_role = aws.iam.Role(f"{config.project_name}-task-role",
        assume_role_policy=json.dumps({
            "Statement": [{
                "Action": "sts:AssumeRole",
                "Effect": "Allow",
                "Principal": {
                    "Service": "ecs-tasks.amazonaws.com",
                },
            }],
            "Version": "2012-10-17",
        }),
        inline_policies=[
            aws.iam.RoleInlinePolicyArgs(
                name="ExecuteCommand",
                policy=json.dumps({
                    "Version": "2012-10-17",
                    "Statement": [
                        {
                            "Action": [
                                "ssmmessages:CreateControlChannel",
                                "ssmmessages:OpenControlChannel",
                                "ssmmessages:CreateDataChannel",
                                "ssmmessages:OpenDataChannel",
                            ],
                            "Effect": "Allow",
                            "Resource": "*",
                        },
                        {
                            "Action": [
                                "logs:CreateLogStream",
                                "logs:DescribeLogGroups",
                                "logs:DescribeLogStreams",
                                "logs:PutLogEvents",
                            ],
                            "Effect": "Allow",
                            "Resource": "*",
                        },
                    ],
                }),
            ),
            aws.iam.RoleInlinePolicyArgs(
                name="DenyIAM",
                policy=json.dumps({
                    "Version": "2012-10-17",
                    "Statement": [{
                        "Action": "iam:*",
                        "Effect": "Deny",
                        "Resource": "*",
                    }],
                }),
            ),
            aws.iam.RoleInlinePolicyArgs(
                name="BedrockS3SQSAccess",
                policy=json.dumps({
                    "Version": "2012-10-17",
                    "Statement": [
                        # S3
                        {
                            "Effect": "Allow",
                            "Action": [
                                "s3:*",
                                "s3-object-lambda:*"
                            ],
                            "Resource": "*"
                        },
                        # SQS
                        {
                            "Effect": "Allow",
                            "Action": [
                                "sqs:StartMessageMoveTask",
                                "sqs:DeleteMessage",
                                "sqs:GetQueueUrl",
                                "sqs:ListDeadLetterSourceQueues",
                                "sqs:ListMessageMoveTasks",
                                "sqs:PurgeQueue",
                                "sqs:ReceiveMessage",
                                "sqs:GetQueueAttributes",
                                "sqs:ListQueueTags"
                            ],
                            "Resource": "arn:aws:sqs:us-east-1:673991670544:ai-med-dev-queue-63cb463"
                        },
                        {
                            "Effect": "Allow",
                            "Action": "sqs:ListQueues",
                            "Resource": "*"
                        },
                        # Bedrock
                        {
                            "Effect": "Allow",
                            "Action": [
                                "bedrock:*"
                            ],
                            "Resource": "*"
                        },
                        {
                            "Effect": "Allow",
                            "Action": [
                                "kms:DescribeKey"
                            ],
                            "Resource": "arn:*:kms:*:*:key/*"
                        },
                        {
                            "Effect": "Allow",
                            "Action": [
                                "iam:ListRoles",
                                "ec2:DescribeVpcs",
                                "ec2:DescribeSubnets",
                                "ec2:DescribeSecurityGroups"
                            ],
                            "Resource": "*"
                        },
                        {
                            "Effect": "Allow",
                            "Action": [
                                "iam:PassRole"
                            ],
                            "Resource": "arn:aws:iam::*:role/*AmazonBedrock*",
                            "Condition": {
                                "StringEquals": {
                                    "iam:PassedToService": "bedrock.amazonaws.com"
                                }
                            }
                        }   
                    ]
                })
            ),            
        ])
    return task_role