# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0

Description: This template manages an MWAA environment and related components

Parameters:
  AirflowBucketName:
    Description: The bucket name of the Airflow environment which stores the DAGs, plugins, etc.
    Type: String

  EnvironmentName:
    Description: Name of the MWAA Environment
    Type: String
    Default: MWAA-Batch-Compute-Environment

  MWAAVersion:
    Description: Version of MWAA to create
    Type: String
    Default: 2.4.3

  VpcCIDR:
    Description: IP range (CIDR notation) for MWAA VPC
    Type: String
    Default: 10.192.0.0/16

  PublicSubnet1CIDR:
    Description: IP range (CIDR notation) for the public subnet in the first Availability Zone
    Type: String
    Default: 10.192.10.0/24

  PublicSubnet2CIDR:
    Description: IP range (CIDR notation) for the public subnet in the second Availability Zone
    Type: String
    Default: 10.192.11.0/24

  PrivateSubnet1CIDR:
    Description: IP range (CIDR notation) for the private subnet in the first Availability Zone
    Type: String
    Default: 10.192.20.0/24

  PrivateSubnet2CIDR:
    Description: IP range (CIDR notation) for the private subnet in the second Availability Zone
    Type: String
    Default: 10.192.21.0/24

Resources:
  VPC:
    Type: AWS::EC2::VPC
    Properties:
      CidrBlock: !Ref VpcCIDR
      EnableDnsSupport: true
      EnableDnsHostnames: true
      Tags:
        - Key: Name
          Value: !Ref EnvironmentName

  InternetGateway:
    Type: AWS::EC2::InternetGateway
    Properties:
      Tags:
        - Key: Name
          Value: !Ref EnvironmentName

  InternetGatewayAttachment:
    Type: AWS::EC2::VPCGatewayAttachment
    Properties:
      InternetGatewayId: !Ref InternetGateway
      VpcId: !Ref VPC

  PublicSubnet1:
    Type: AWS::EC2::Subnet
    Properties:
      VpcId: !Ref VPC
      AvailabilityZone: !Select [ 0, !GetAZs '' ]
      CidrBlock: !Ref PublicSubnet1CIDR
      MapPublicIpOnLaunch: true
      Tags:
        - Key: Name
          Value: !Sub ${EnvironmentName} Public Subnet (AZ1)

  PublicSubnet2:
    Type: AWS::EC2::Subnet
    Properties:
      VpcId: !Ref VPC
      AvailabilityZone: !Select [ 1, !GetAZs  '' ]
      CidrBlock: !Ref PublicSubnet2CIDR
      MapPublicIpOnLaunch: true
      Tags:
        - Key: Name
          Value: !Sub ${EnvironmentName} Public Subnet (AZ2)

  PrivateSubnet1:
    Type: AWS::EC2::Subnet
    Properties:
      VpcId: !Ref VPC
      AvailabilityZone: !Select [ 0, !GetAZs  '' ]
      CidrBlock: !Ref PrivateSubnet1CIDR
      MapPublicIpOnLaunch: false
      Tags:
        - Key: Name
          Value: !Sub ${EnvironmentName} Private Subnet (AZ1)

  PrivateSubnet2:
    Type: AWS::EC2::Subnet
    Properties:
      VpcId: !Ref VPC
      AvailabilityZone: !Select [ 1, !GetAZs  '' ]
      CidrBlock: !Ref PrivateSubnet2CIDR
      MapPublicIpOnLaunch: false
      Tags:
        - Key: Name
          Value: !Sub ${EnvironmentName} Private Subnet (AZ2)

  NatGateway1EIP:
    Type: AWS::EC2::EIP
    DependsOn: InternetGatewayAttachment
    Properties:
      Domain: vpc

  NatGateway2EIP:
    Type: AWS::EC2::EIP
    DependsOn: InternetGatewayAttachment
    Properties:
      Domain: vpc

  NatGateway1:
    Type: AWS::EC2::NatGateway
    Properties:
      AllocationId: !GetAtt NatGateway1EIP.AllocationId
      SubnetId: !Ref PublicSubnet1

  NatGateway2:
    Type: AWS::EC2::NatGateway
    Properties:
      AllocationId: !GetAtt NatGateway2EIP.AllocationId
      SubnetId: !Ref PublicSubnet2

  PublicRouteTable:
    Type: AWS::EC2::RouteTable
    Properties:
      VpcId: !Ref VPC
      Tags:
        - Key: Name
          Value: !Sub ${EnvironmentName} Public Routes

  DefaultPublicRoute:
    Type: AWS::EC2::Route
    DependsOn: InternetGatewayAttachment
    Properties:
      RouteTableId: !Ref PublicRouteTable
      DestinationCidrBlock: 0.0.0.0/0
      GatewayId: !Ref InternetGateway

  PublicSubnet1RouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      RouteTableId: !Ref PublicRouteTable
      SubnetId: !Ref PublicSubnet1

  PublicSubnet2RouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      RouteTableId: !Ref PublicRouteTable
      SubnetId: !Ref PublicSubnet2


  PrivateRouteTable1:
    Type: AWS::EC2::RouteTable
    Properties:
      VpcId: !Ref VPC
      Tags:
        - Key: Name
          Value: !Sub ${EnvironmentName} Private Routes (AZ1)

  DefaultPrivateRoute1:
    Type: AWS::EC2::Route
    Properties:
      RouteTableId: !Ref PrivateRouteTable1
      DestinationCidrBlock: 0.0.0.0/0
      NatGatewayId: !Ref NatGateway1

  PrivateSubnet1RouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      RouteTableId: !Ref PrivateRouteTable1
      SubnetId: !Ref PrivateSubnet1

  PrivateRouteTable2:
    Type: AWS::EC2::RouteTable
    Properties:
      VpcId: !Ref VPC
      Tags:
        - Key: Name
          Value: !Sub ${EnvironmentName} Private Routes (AZ2)

  DefaultPrivateRoute2:
    Type: AWS::EC2::Route
    Properties:
      RouteTableId: !Ref PrivateRouteTable2
      DestinationCidrBlock: 0.0.0.0/0
      NatGatewayId: !Ref NatGateway2

  PrivateSubnet2RouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      RouteTableId: !Ref PrivateRouteTable2
      SubnetId: !Ref PrivateSubnet2

  MwaaVpcSecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupName: mwaa-sg
      GroupDescription: Security Group for Amazon MWAA Environment
      VpcId: !Ref VPC
  
  MwaaVpcSecurityGroupInboundRule:
    Type: AWS::EC2::SecurityGroupIngress
    Properties:
      IpProtocol: '-1'
      SourceSecurityGroupId: !GetAtt MwaaVpcSecurityGroup.GroupId
      GroupId: !GetAtt MwaaVpcSecurityGroup.GroupId

  MwaaRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
              - airflow.amazonaws.com
              - airflow-env.amazonaws.com
            Action:
              - sts:AssumeRole
      Policies:
        - PolicyName: MwaaMinimalExecutionRolePolicy
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action: airflow:PublishMetrics
                Resource: !Sub arn:${AWS::Partition}:airflow:${AWS::Region}:${AWS::AccountId}:environment/${EnvironmentName}
              - Effect: Allow
                Action:
                  - s3:GetObject*
                  - s3:GetBucket*
                  - s3:List*
                Resource:
                  - !Sub arn:${AWS::Partition}:s3:::${AirflowBucketName}
                  - !Sub arn:${AWS::Partition}:s3:::${AirflowBucketName}/*
              - Effect: Allow
                Action:
                  - logs:CreateLogStream
                  - logs:CreateLogGroup
                  - logs:PutLogEvents
                  - logs:GetLogEvents
                  - logs:GetLogRecord
                  - logs:GetLogGroupFields
                  - logs:GetQueryResults
                Resource: 
                  - !Sub arn:${AWS::Partition}:logs:${AWS::Region}:${AWS::AccountId}:log-group:*
              - Effect: Allow
                Action: logs:DescribeLogGroups
                Resource: '*'
              - Effect: Allow
                Action: 
                  - s3:GetAccountPublicAccessBlock
                  - "s3:List*"
                  - s3:GetEncryptionConfiguration
                Resource: '*'
              - Effect: Allow
                Action: cloudwatch:PutMetricData
                Resource: '*'
              - Effect: Allow
                Action:
                  - sqs:ChangeMessageVisibility
                  - sqs:DeleteMessage
                  - sqs:GetQueueAttributes
                  - sqs:GetQueueUrl
                  - sqs:ReceiveMessage
                  - sqs:SendMessage
                Resource: !Sub arn:${AWS::Partition}:sqs:${AWS::Region}:*:airflow-celery-*
              - Effect: Allow
                Action:
                  - kms:Decrypt
                  - kms:DescribeKey
                  - kms:GenerateDataKey*
                  - kms:Encrypt
                NotResource: !Sub arn:${AWS::Partition}:kms:*:${AWS::AccountId}:key/*
                Condition:
                  StringLike:
                    'kms:ViaService': !Sub sqs.${AWS::Region}.amazonaws.com
                # Permissions needed for the ECSOperator and to scale the ECS Cluster
              - Effect: Allow
                Action:
                  - ecs:RunTask
                  - ecs:StopTask 
                  - ecs:DescribeTasks
                  - ecs:ListClusters
                  - ecs:DescribeClusters 
                  - ecs:ListServices
                  - ecs:DescribeServices
                  - ecs:ListContainerInstances
                  - ecs:DescribeContainerInstances
                  - ecs:DescribeTaskDefinition
                  - ecs:ListTaskDefinitions
                  - ecs:ListTagsForResource
                  - autoscaling:SetDesiredCapacity
                  - autoscaling:DescribeAutoScalingGroups
                  - autoscaling:UpdateAutoScalingGroup
                  - ec2:DescribeInstances
                  - ec2:DescribeInstanceTypes
                Resource: '*'
              - Effect: Allow
                Action: iam:PassRole
                Resource: '*'
                Condition:
                  StringLike:
                    'iam:PassedToService': 'ecs-tasks.amazonaws.com'

  AirflowEnvironment:
    Type: AWS::MWAA::Environment
    DependsOn:
      - NatGateway1
      - NatGateway2
    Properties:
      Name: !Ref EnvironmentName
      AirflowVersion: !Ref MWAAVersion
      NetworkConfiguration:
        SubnetIds:
          - !Ref PrivateSubnet1
          - !Ref PrivateSubnet2
        SecurityGroupIds:
          - !Ref MwaaVpcSecurityGroup
      LoggingConfiguration:
        DagProcessingLogs:
          Enabled: true
          LogLevel: INFO
        SchedulerLogs:
          Enabled: true
          LogLevel: INFO
        TaskLogs:
          Enabled: true
          LogLevel: INFO
        WebserverLogs:
          Enabled: true
          LogLevel: INFO
        WorkerLogs:
          Enabled: true
          LogLevel: INFO
      SourceBucketArn: !Sub arn:${AWS::Partition}:s3:::${AirflowBucketName}
      ExecutionRoleArn: !GetAtt MwaaRole.Arn
      MinWorkers: 1
      MaxWorkers: 10
      DagS3Path: dags
      EnvironmentClass: mw1.small
      WebserverAccessMode: PUBLIC_ONLY
      AirflowConfigurationOptions:
        'core.dag_run_conf_overrides_params' : 'True'

Outputs:
  MwaaWebserverUrl:
    Description: The URL of your Apache Airflow UI
    Value: !GetAtt AirflowEnvironment.WebserverUrl