Parameters:
  RuleGroupName:
    Type: String
    Default: 'AutoUpdating-SFTP-FQDN'
  RuleGroupFQDN:
    Type: String
    Default: www.aws.com
    Description: "Type the FQDN of the SFTP endpoint"
    MinLength: 6
    MaxLength: 255
    AllowedPattern: ^(?!:\/\/)(?=.{1,255}$)((.{1,63}\.){1,127}(?![0-9]*$)[a-z0-9-]+\.?)$
    ConstraintDescription: must contain a valid Fully Qualified Domain Name (FQDN)

Resources:
  StatefulRulegroup:
    Type: 'AWS::NetworkFirewall::RuleGroup'
    Properties:
      Capacity: 100
      RuleGroupName: !Sub "${AWS::StackName}-${RuleGroupName}"
      Description: This Rule group is used to restrict SFTP access to an FQDN by IP addresses fetched via Lambda
      Type: STATEFUL
      RuleGroup:
        RuleVariables:
            IPSets:
              "SFTPFQDN":
                Definition: ["127.0.0.1"]
        RulesSource:
          RulesString: !Sub |
            # This will be updated automatically by the Lambda
      Tags: 
            - Key: "ProjectName"
              Value: "SFTP-FQDN"
  LambdaExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
        - Effect: Allow
          Principal:
            Service:
            - lambda.amazonaws.com
          Action:
          - sts:AssumeRole
      Path: "/"
      Policies:
      - PolicyName: LambdaLogs
        PolicyDocument:
          Version: '2012-10-17'
          Statement:
          - Effect: Allow
            Action:
              - 'logs:CreateLogStream'
              - 'logs:PutLogEvents'
            Resource: !Sub 'arn:aws:logs:${AWS::Region}:${AWS::AccountId}:*'
      - PolicyName: NetworkFirewall
        PolicyDocument:
          Version: '2012-10-17'
          Statement:
          - Effect: Allow
            Action:
            - 'network-firewall:*'
            Resource: 
              - !GetAtt StatefulRulegroup.RuleGroupArn
      Tags: 
          - Key: "ProjectName"
            Value: "SFTP-FQDN"
  LambdaLogGroup:
    Type: AWS::Logs::LogGroup
    Properties:
      LogGroupName: !Sub "/aws/lambda/${LambdaFunction}"
    DeletionPolicy: Retain
  ScheduledRule: 
    Type: AWS::Events::Rule
    Properties: 
      Description: "SFTP-FQDNTrigger"
      ScheduleExpression: "cron(0/5 * * * ? *)"
      State: "ENABLED"
      Targets: 
        - 
          Arn: 
            Fn::GetAtt: 
              - "LambdaFunction"
              - "Arn"
          Id: "TargetFunctionV1"
  PermissionForEventsToInvokeLambda: 
    Type: AWS::Lambda::Permission
    Properties: 
      FunctionName: !Ref "LambdaFunction"
      Action: "lambda:InvokeFunction"
      Principal: "events.amazonaws.com"
      SourceArn: 
        Fn::GetAtt: 
          - "ScheduledRule"
          - "Arn"
  LambdaInvoke:
    Type: AWS::CloudFormation::CustomResource
    Version: "1.0"
    Properties:
      ServiceToken: !GetAtt LambdaFunction.Arn
  LambdaFunction:
    Type: AWS::Lambda::Function
    Properties:
      Role: !GetAtt LambdaExecutionRole.Arn
      Runtime: nodejs16.x
      Handler: index.handler
      Timeout: 60
      Description: Used to fetch IPs for the given FQDN and update the associated RuleGroup
      Tags: 
        - 
          Key: "ProjectName"
          Value: "SFTP-FQDN"
      Code:
        ZipFile: !Sub |
          const AWS = require("aws-sdk");
          var response = require('cfn-response');
          const dnsPromises = require('dns').promises;

          const networkfirewall = new AWS.NetworkFirewall();

          const getAddresses = async function(fqdn){
            let res = await dnsPromises.resolve4(fqdn);
            return res.map((line)=> {return line + "/32"});
          };

          const updateRules = async function (ruleGroup, fqdn, addresses) {
            let params = ruleGroup;
            params.RuleGroupName = params.RuleGroupResponse.RuleGroupName;
            params.Description = params.RuleGroupResponse.Description;
            params.Type = params.RuleGroupResponse.Type;
            delete params.Capacity;
            delete params.RuleGroupResponse;
            
            let rulesString = "# Last autofetched by Lambda: " + new Date().toUTCString() + "\n";
            rulesString += "# Fetched addresses for: " + fqdn + " stored as Variable: $SFTPFQDN\n";
            addresses.forEach (address => {
              rulesString += "# " + address + "\n";
            });
            rulesString += 'pass tcp any any -> $SFTPFQDN 22 (msg:"Allow access to ' + fqdn + '"; sid:1001;)';
            
            params.RuleGroup.RulesSource.RulesString = rulesString;
            params.RuleGroup.RuleVariables.IPSets.SFTPFQDN.Definition = addresses;

            console.log("Updating rules...");
            let res = await networkfirewall.updateRuleGroup(params).promise();
            if (res) {
              console.log("Updated '" + params.RuleGroupName + "'.");
            } else {
              console.log("Error updating the rules for '" + params.RuleGroupName + "'...");
            }
            
            return;
          };

          exports.handler = async (event, context) => {
            if (event.RequestType == "Delete") {
              await response.send(event, context, "SUCCESS");
              return;
            }

            var rg = {Type: "STATEFUL", RuleGroupArn: '${StatefulRulegroup.RuleGroupArn}'};
            const fqdn = "${RuleGroupFQDN}";
            
            let addresses = await getAddresses(fqdn);
            if (addresses) {
              console.log("Searching Rule Groups for " + rg.RuleGroupArn + "...");
              let res = await networkfirewall.describeRuleGroup(rg).promise();
              if (res.RuleGroupResponse) {
                console.log("Found matching Rule Group...");
                await updateRules(res, fqdn, addresses);
                if (event.ResponseURL) await response.send(event, context, response.SUCCESS);
              } else {
                console.log("ERROR: No matching Rule Group found...");
                if (event.ResponseURL) await response.send(event, context, response.FAILED);
              }  
            } else {
              console.log("Could not resolve addresses for fqdn: " + fqdn);
              if (event.ResponseURL) await response.send(event, context, response.FAILED);
            }
            
            return;
          };