#Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.

#THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
#FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
#COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
#IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
#CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

AWSTemplateFormatVersion: "2010-09-09"
Description: "AWS Network Firewall Demo using distributed model."

Metadata:
  "AWS::CloudFormation::Interface":
    ParameterGroups:
      - Label:
          default: "VPC Parameters"
        Parameters: 
          - AvailabilityZoneSelection    
      - Label:
          default: "EC2 Parameters"
        Parameters: 
          - LatestAmiId

Parameters:
  AvailabilityZoneSelection:
    Description: Availability Zone
    Type: AWS::EC2::AvailabilityZone::Name
    Default: us-east-1a


  LatestAmiId:
    Description: Latest EC2 AMI from Systems Manager Parameter Store
    Type: 'AWS::SSM::Parameter::Value<AWS::EC2::Image::Id>'
    Default: '/aws/service/ami-amazon-linux-latest/amzn2-ami-hvm-x86_64-gp2'
      
Resources:

# VPC:
  DemoVPC:
    Type: AWS::EC2::VPC
    Properties:
      CidrBlock: "10.1.0.0/16"
      EnableDnsSupport: "true"
      EnableDnsHostnames: "true"
      InstanceTenancy: default
      Tags:
        - Key: Name
          Value: !Sub "${AWS::StackName}-DemoVPC"

# Internet Gateway:
  InternetGateway:
    Type: AWS::EC2::InternetGateway
    Properties:
      Tags:
        - Key: Name
          Value: !Sub "${AWS::StackName}-DemoVPC-igw"

  AttachInternetGateway:
    Type: AWS::EC2::VPCGatewayAttachment
    Properties:
      VpcId:
        !Ref DemoVPC
      InternetGatewayId:
        !Ref InternetGateway

# NAT Gateway:
  NatGwEip:
    Type: AWS::EC2::EIP
    Properties:
      Domain: vpc


  NatGw:
    Type: AWS::EC2::NatGateway
    DependsOn:
      - NatGwEip
      - PublicSubnet
    Properties:
      AllocationId: !GetAtt
        - NatGwEip
        - AllocationId
      SubnetId: !Ref PublicSubnet
      Tags:
        - Key: Name
          Value: !Sub "${AWS::StackName}-DemoVPC-natgw"


# Private Subnets for Test Instances:
  PrivateSubnet:
    Type: AWS::EC2::Subnet
    Properties:
      VpcId:
        Ref: DemoVPC
      CidrBlock: "10.1.0.0/24"
      AvailabilityZone:
        Ref: AvailabilityZoneSelection
      MapPublicIpOnLaunch: false
      Tags:
        - Key: Name
          Value: !Sub "${AWS::StackName}-private-subnet"


# Public Subnets for NAT GWs:
  PublicSubnet:
    Type: AWS::EC2::Subnet
    Properties:
      VpcId:
        Ref: DemoVPC
      CidrBlock: "10.1.1.0/24"
      AvailabilityZone:
        Ref: AvailabilityZoneSelection
      MapPublicIpOnLaunch: true
      Tags:
        - Key: Name
          Value: !Sub "${AWS::StackName}-public-subnet"


# Firewall Subnets for firewall endpoints:
  FirewallSubnet:
    Type: AWS::EC2::Subnet
    Properties:
      VpcId:
        Ref: DemoVPC
      CidrBlock: "10.1.16.0/28"
      AvailabilityZone:
        Ref: AvailabilityZoneSelection
      MapPublicIpOnLaunch: false
      Tags:
        - Key: Name
          Value: !Sub "${AWS::StackName}-firewall-subnet"

# AWS PrivateLink interface endpoint for services:
  DemoVPCEndpointSecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties:
        GroupDescription: Allow instances to get to SSM Systems Manager
        VpcId: !Ref DemoVPC
        SecurityGroupIngress:
        - IpProtocol: tcp
          FromPort: 443
          ToPort: 443
          CidrIp: 10.1.0.0/16
        Tags:
          - Key: Name
            Value: !Sub "${AWS::StackName}-vpce-sg"          

  DemoVPCSSMEndpoint:
    Type: AWS::EC2::VPCEndpoint
    Properties: 
        PrivateDnsEnabled: true
        SecurityGroupIds: 
          - !Ref DemoVPCEndpointSecurityGroup
        ServiceName: !Sub "com.amazonaws.${AWS::Region}.ssm"
        SubnetIds: 
          - !Ref PublicSubnet
        VpcEndpointType: Interface
        VpcId: !Ref DemoVPC

  DemoVPCEC2MessagesEndpoint:
    Type: AWS::EC2::VPCEndpoint
    Properties: 
        PrivateDnsEnabled: true
        SecurityGroupIds: 
          - !Ref DemoVPCEndpointSecurityGroup
        ServiceName: !Sub "com.amazonaws.${AWS::Region}.ec2messages"
        SubnetIds: 
          - !Ref PublicSubnet
        VpcEndpointType: Interface
        VpcId: !Ref DemoVPC

  DemoVPCSSMMessagesEndpoint:
    Type: AWS::EC2::VPCEndpoint
    Properties: 
        PrivateDnsEnabled: true
        SecurityGroupIds: 
          - !Ref DemoVPCEndpointSecurityGroup
        ServiceName: !Sub "com.amazonaws.${AWS::Region}.ssmmessages"
        SubnetIds: 
          - !Ref PublicSubnet
        VpcEndpointType: Interface
        VpcId: !Ref DemoVPC

# SSM Role:
  SubnetRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: !Sub "${AWS::StackName}-subnet-role"
      Path: "/"
      ManagedPolicyArns:
        - "arn:aws:iam::aws:policy/service-role/AmazonEC2RoleforSSM"
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - ec2.amazonaws.com
            Action:
              - sts:AssumeRole

  SubnetInstanceProfile:
    Type: AWS::IAM::InstanceProfile
    Properties:
      Path: "/"
      Roles:
        - !Ref SubnetRole

# Fn::GetAtt for Firewall do not return VPCE Id in ordered format.
# For more details refer to: https://github.com/aws-cloudformation/aws-cloudformation-resource-providers-networkfirewall/issues/15
# Until the bug is fixed we have to rely on custom resource to retrieve AZ specific VPCE Id.

# Lambda Role:
  LambdaExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: !Sub "${AWS::StackName}-lambda-role"
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - lambda.amazonaws.com
            Action:
              - sts:AssumeRole
      Path: /
      Policies:
        - PolicyName: root
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action:
                  - logs:CreateLogStream
                  - logs:PutLogEvents
                Resource: !GetAtt RetrieveVpcIdLogGroup.Arn
              - Effect: Allow
                Action:
                  - network-firewall:DescribeFirewall
                Resource: "*"

# Retrieve VpceId Lambda Custom Resource:
  RetrieveVpcIdLogGroup:
    Type: AWS::Logs::LogGroup
    Properties:
        LogGroupName: !Sub /aws/lambda/${AWS::StackName}-retrieve-vpceid
        RetentionInDays: 1

  RetrieveVpceId:
    Type: AWS::Lambda::Function
    DependsOn: RetrieveVpcIdLogGroup
    Properties:
      FunctionName: !Sub ${AWS::StackName}-retrieve-vpceid
      Handler: "index.handler"
      Role: !GetAtt
        - LambdaExecutionRole
        - Arn
      Code:
        ZipFile: |
          import boto3
          import cfnresponse
          import json
          import logging

          def handler(event, context):
              logger = logging.getLogger()
              logger.setLevel(logging.INFO)
              responseData = {}
              responseStatus = cfnresponse.FAILED
              logger.info('Received event: {}'.format(json.dumps(event)))
              if event["RequestType"] == "Delete":
                  responseStatus = cfnresponse.SUCCESS
                  cfnresponse.send(event, context, responseStatus, responseData)
              if event["RequestType"] == "Create":
                  try:
                      Az = event["ResourceProperties"]["Az"]
                      FwArn = event["ResourceProperties"]["FwArn"]
                  except Exception as e:
                      logger.info('AZ retrieval failure: {}'.format(e))
                  try:
                      nfw = boto3.client('network-firewall')
                  except Exception as e:
                      logger.info('boto3.client failure: {}'.format(e))
                  try:
                      NfwResponse=nfw.describe_firewall(FirewallArn=FwArn)
                      VpceId = NfwResponse['FirewallStatus']['SyncStates'][Az]['Attachment']['EndpointId']


                  except Exception as e:
                      logger.info('ec2.describe_firewall failure: {}'.format(e))

                  responseData['FwVpceId'] = VpceId
                  responseStatus = cfnresponse.SUCCESS
                  cfnresponse.send(event, context, responseStatus, responseData)
      Runtime: python3.7
      Timeout: 30

  FirewallVpceIds:
    Type: Custom::DescribeVpcEndpoints
    Properties:
      ServiceToken: !GetAtt RetrieveVpceId.Arn
      Az: !Ref AvailabilityZoneSelection
      FwArn: !Ref DemoVPCFirewall

# Testing Security Group:
  SubnetSecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupDescription: "ICMP acess from 10.0.0.0/8"
      GroupName: !Sub "${AWS::StackName}-test-instance-sec-group"
      VpcId: !Ref DemoVPC
      SecurityGroupIngress:
        - IpProtocol: icmp
          CidrIp: 10.0.0.0/8
          FromPort: "-1"
          ToPort: "-1"
      Tags:
        - Key: Name
          Value: !Sub "${AWS::StackName}-test-sg"                

# Test Instances:
  TestInstance:
    Type: 'AWS::EC2::Instance'
    Properties:
      ImageId: !Ref LatestAmiId
      SubnetId: !Ref PrivateSubnet
      InstanceType: t2.micro
      SecurityGroupIds:
        - !Ref SubnetSecurityGroup
      IamInstanceProfile: !Ref SubnetInstanceProfile
      Tags:
        - Key: Name
          Value: !Sub "${AWS::StackName}-test-instance"


# AWS Network Firewall:
  DemoVPCFirewall:
    Type: AWS::NetworkFirewall::Firewall
    Properties:
      FirewallName: !Sub "${AWS::StackName}-firewall"
      FirewallPolicyArn: !Ref EgressFirewallPolicy
      VpcId: !Ref DemoVPC
      SubnetMappings:
        - SubnetId: !Ref FirewallSubnet
      Tags:
        - Key: Name
          Value: !Sub "${AWS::StackName}-firewall"
  
  WICARAlertStatefulRuleGroup:
    Type: 'AWS::NetworkFirewall::RuleGroup'
    Properties:
      RuleGroupName: SampleStatefulRulegroupName
      Type: STATEFUL
      RuleGroup:
        RulesSource:
          RulesString: >
            drop http $HOME_NET any -> $EXTERNAL_NET 80 (msg:"MALWARE custom solution"; flow: to_server,established; classtype:trojan-activity; sid:2002001; content:"/data/js_crypto_miner.html";http_uri; rev:1;)
            
            drop http $HOME_NET any -> $EXTERNAL_NET 80 (msg:"MALWARE custom solution"; flow: to_server,established; classtype:trojan-activity; sid:2002002; content:"/data/java_jre17_exec.html";http_uri; rev:1;)
      Capacity: 100
      Description: Custom Rulegroup
      Tags:
        - Key: Name
          Value: !Sub "Custom-alert-${AWS::StackName}"



          

  EgressFirewallPolicy:
    Type: AWS::NetworkFirewall::FirewallPolicy
    Properties:
      FirewallPolicyName: !Sub "${AWS::StackName}-firewall-policy"
      FirewallPolicy:
        StatelessDefaultActions:
          - 'aws:forward_to_sfe'
        StatelessFragmentDefaultActions:
          - 'aws:forward_to_sfe'


      Tags:
        - Key: Name
          Value: !Sub "${AWS::StackName}-firewall-policy"

  DemoVPCFirewallLogFlowGroup:
    Type: AWS::Logs::LogGroup
    Properties:
      LogGroupName: !Sub "/${AWS::StackName}/anfw/flow"

  DemoVPCFirewallLogAlertGroup:
    Type: AWS::Logs::LogGroup
    Properties:
      LogGroupName: !Sub "/${AWS::StackName}/anfw/alert"

  DemoVPCFirewallLog:
    Type: AWS::NetworkFirewall::LoggingConfiguration
    Properties:
      FirewallArn: !Ref DemoVPCFirewall
      LoggingConfiguration:
        LogDestinationConfigs:
          - LogType: FLOW
            LogDestinationType: CloudWatchLogs
            LogDestination:
              logGroup: !Sub "/${AWS::StackName}/anfw/flow"
          - LogType: ALERT
            LogDestinationType: CloudWatchLogs
            LogDestination:
              logGroup: !Sub "/${AWS::StackName}/anfw/alert"

# Private Route Tables:
  PrivateRtb:
    Type: AWS::EC2::RouteTable
    Properties:
      VpcId: !Ref DemoVPC
      Tags:
        - Key: Name
          Value: !Sub "${AWS::StackName}-private-route-table"

  PrivateRtbAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    DependsOn: PrivateSubnet
    Properties:
      RouteTableId: !Ref PrivateRtb
      SubnetId: !Ref PrivateSubnet

  PrivateRtbDefaultRoute:
    Type: AWS::EC2::Route
    DependsOn: NatGw
    Properties:
      DestinationCidrBlock: "0.0.0.0/0"
      NatGatewayId: !Ref NatGw
      RouteTableId: !Ref PrivateRtb


# Public Route Tables:
  PublicRtb:
    Type: AWS::EC2::RouteTable
    Properties:
      VpcId: !Ref DemoVPC
      Tags:
        - Key: Name
          Value: !Sub "${AWS::StackName}-public-route-table"

  PublicRtbAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    DependsOn: PublicSubnet
    Properties:
      RouteTableId: !Ref PublicRtb
      SubnetId: !Ref PublicSubnet

  PublicRtbDefaultRoute:
    Type: AWS::EC2::Route
    DependsOn: DemoVPCFirewall
    Properties:
      DestinationCidrBlock: "0.0.0.0/0"
      VpcEndpointId: !GetAtt FirewallVpceIds.FwVpceId
      RouteTableId: !Ref PublicRtb


# Firewall Route Tables:
  FirewallRtb:
    Type: AWS::EC2::RouteTable
    Properties:
      VpcId: !Ref DemoVPC
      Tags:
        - Key: Name
          Value: !Sub "${AWS::StackName}-firewall-route-table"

  FirewallRtbAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    DependsOn: FirewallSubnet
    Properties:
      RouteTableId: !Ref FirewallRtb
      SubnetId: !Ref FirewallSubnet
      
  FirewallRtbDefaultRoute:
    Type: AWS::EC2::Route
    DependsOn: InternetGateway
    Properties:
      DestinationCidrBlock: "0.0.0.0/0"
      GatewayId: !Ref InternetGateway
      RouteTableId: !Ref FirewallRtb

# Ingress Route Table:
  IngressRtb:
    Type: AWS::EC2::RouteTable
    Properties:
      VpcId: !Ref DemoVPC
      Tags:
        - Key: Name
          Value: !Sub "${AWS::StackName}-ingress-route-table"
  
  IngressRtbAssociation:
    Type: AWS::EC2::GatewayRouteTableAssociation
    DependsOn: InternetGateway
    Properties:
      RouteTableId: !Ref IngressRtb
      GatewayId: !Ref InternetGateway
      
  IngressRtbPublicSubnetRoute:
    Type: AWS::EC2::Route
    DependsOn: DemoVPCFirewall
    Properties:
      DestinationCidrBlock: "10.1.1.0/24"
      VpcEndpointId: !GetAtt FirewallVpceIds.FwVpceId
      RouteTableId: !Ref IngressRtb
