# AWS DISCLAIMER
# ---
# The following files are provided by AWS Professional Services describe the process to create an SSM with description.

# These are non-production ready and are to be used for testing purposes.

# These files is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
# OR CONDITIONS OF ANY KIND, either express or implied. See the License
# for the specific language governing permissions and limitations under the License.

# (c) 2020 Amazon Web Services, Inc. or its affiliates. All Rights Reserved.
# This AWS Content is provided subject to the terms of the AWS Customer Agreement available at
# http://aws.amazon.com/agreement or other written agreement between Customer and Amazon Web Services, Inc.​
AWSTemplateFormatVersion: "2010-09-09"
Description: "This cloudformation template creates a FSx for Windows File System and enabled shadow copying with a custom resource"
Parameters:
  LatestWindowsAmiId:
    Type: AWS::SSM::Parameter::Value<AWS::EC2::Image::Id>
    Default: /aws/service/ami-windows-latest/Windows_Server-2019-English-Full-Base
  VpcId:
    Type: AWS::EC2::VPC::Id
    Description: Enter the security group ID to attach to the FSx for Windows file system
  SubnetId:
    Type: AWS::EC2::Subnet::Id
    Description: Enter the subnet ID that the FSx for Windows file system should be placed in
  ActiveDirectoryId:
    Type: String
    Description: Enter the Active Directory ID that the FSx for Windows file system will be associated with
  ActiveDirectoryName:
    Type: String
    Description: Enter the Active Directory ID that the FSx for Windows file system will be associated with
  ActiveDirectoryDnsAddresses:
    Type: List<String>
    Description: Enter the IP Address of the Active Directory DNS servers
  ManagementInstanceAvailabilityZoneName:
    Type: AWS::EC2::AvailabilityZone::Name
    Description: Enter the name of the Availability Zone that EC2 instance that will be used to manage the FSx for Windows file system will be placed in
  ActiveDirectoryCredentialsParameterName:
    Type: String
    Description: Enter the name of the Systems Manager Parameter Store parameter name containing the Active Directory credentials JSON string
Resources:
  FSxSecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupDescription: Allows traffic to and from FSx file system
      VpcId: !Ref VpcId
  FSxSecurityGroupIngressUdpDns:
    Type: AWS::EC2::SecurityGroupIngress
    Properties:
      FromPort: 53
      ToPort: 53
      IpProtocol: udp
      GroupId: !GetAtt FSxSecurityGroup.GroupId
      SourceSecurityGroupId: !GetAtt FSxSecurityGroup.GroupId
  FSxSecurityGroupIngressUdp88:
    Type: AWS::EC2::SecurityGroupIngress
    Properties:
      FromPort: 88
      ToPort: 88
      IpProtocol: udp
      GroupId: !GetAtt FSxSecurityGroup.GroupId
      SourceSecurityGroupId: !GetAtt FSxSecurityGroup.GroupId
  FSxSecurityGroupIngressUdp123:
    Type: AWS::EC2::SecurityGroupIngress
    Properties:
      FromPort: 123
      ToPort: 123
      IpProtocol: udp
      GroupId: !GetAtt FSxSecurityGroup.GroupId
      SourceSecurityGroupId: !GetAtt FSxSecurityGroup.GroupId
  FSxSecurityGroupIngressUdpLdap:
    Type: AWS::EC2::SecurityGroupIngress
    Properties:
      FromPort: 389
      ToPort: 389
      IpProtocol: udp
      GroupId: !GetAtt FSxSecurityGroup.GroupId
      SourceSecurityGroupId: !GetAtt FSxSecurityGroup.GroupId
  FSxSecurityGroupIngressUdp464:
    Type: AWS::EC2::SecurityGroupIngress
    Properties:
      FromPort: 464
      ToPort: 464
      IpProtocol: udp
      GroupId: !GetAtt FSxSecurityGroup.GroupId
      SourceSecurityGroupId: !GetAtt FSxSecurityGroup.GroupId
  FSxSecurityGroupIngressTcpDns:
    Type: AWS::EC2::SecurityGroupIngress
    Properties:
      FromPort: 53
      ToPort: 53
      IpProtocol: tcp
      GroupId: !GetAtt FSxSecurityGroup.GroupId
      SourceSecurityGroupId: !GetAtt FSxSecurityGroup.GroupId
  FSxSecurityGroupIngressTcp88:
    Type: AWS::EC2::SecurityGroupIngress
    Properties:
      FromPort: 88
      ToPort: 88
      IpProtocol: tcp
      GroupId: !GetAtt FSxSecurityGroup.GroupId
      SourceSecurityGroupId: !GetAtt FSxSecurityGroup.GroupId
  FSxSecurityGroupIngressTcp123:
    Type: AWS::EC2::SecurityGroupIngress
    Properties:
      FromPort: 123
      ToPort: 123
      IpProtocol: tcp
      GroupId: !GetAtt FSxSecurityGroup.GroupId
      SourceSecurityGroupId: !GetAtt FSxSecurityGroup.GroupId
  FSxSecurityGroupIngressTcp135:
    Type: AWS::EC2::SecurityGroupIngress
    Properties:
      FromPort: 135
      ToPort: 135
      IpProtocol: tcp
      GroupId: !GetAtt FSxSecurityGroup.GroupId
      SourceSecurityGroupId: !GetAtt FSxSecurityGroup.GroupId
  FSxSecurityGroupIngressWinRMHttp:
    Type: AWS::EC2::SecurityGroupIngress
    Properties:
      FromPort: 5985
      ToPort: 5985
      IpProtocol: tcp
      GroupId: !GetAtt FSxSecurityGroup.GroupId
      SourceSecurityGroupId: !GetAtt FSxSecurityGroup.GroupId
  FSxSecurityGroupIngressTcpLdap:
    Type: AWS::EC2::SecurityGroupIngress
    Properties:
      FromPort: 389
      ToPort: 389
      IpProtocol: tcp
      GroupId: !GetAtt FSxSecurityGroup.GroupId
      SourceSecurityGroupId: !GetAtt FSxSecurityGroup.GroupId
  FSxSecurityGroupIngressTcpSmb:
    Type: AWS::EC2::SecurityGroupIngress
    Properties:
      FromPort: 445
      ToPort: 445
      IpProtocol: tcp
      GroupId: !GetAtt FSxSecurityGroup.GroupId
      SourceSecurityGroupId: !GetAtt FSxSecurityGroup.GroupId
  FSxSecurityGroupIngressTcp464:
    Type: AWS::EC2::SecurityGroupIngress
    Properties:
      FromPort: 464
      ToPort: 464
      IpProtocol: tcp
      GroupId: !GetAtt FSxSecurityGroup.GroupId
      SourceSecurityGroupId: !GetAtt FSxSecurityGroup.GroupId
  FSxSecurityGroupIngressTcp636:
    Type: AWS::EC2::SecurityGroupIngress
    Properties:
      FromPort: 636
      ToPort: 636
      IpProtocol: tcp
      GroupId: !GetAtt FSxSecurityGroup.GroupId
      SourceSecurityGroupId: !GetAtt FSxSecurityGroup.GroupId
  FSxSecurityGroupIngressTcp32683269:
    Type: AWS::EC2::SecurityGroupIngress
    Properties:
      FromPort: 3268
      ToPort: 3269
      IpProtocol: tcp
      GroupId: !GetAtt FSxSecurityGroup.GroupId
      SourceSecurityGroupId: !GetAtt FSxSecurityGroup.GroupId
  FSxSecurityGroupIngressTcp9389:
    Type: AWS::EC2::SecurityGroupIngress
    Properties:
      FromPort: 9389
      ToPort: 9389
      IpProtocol: tcp
      GroupId: !GetAtt FSxSecurityGroup.GroupId
      SourceSecurityGroupId: !GetAtt FSxSecurityGroup.GroupId
  FSxSecurityGroupIngressTcpEphemeralPorts:
    Type: AWS::EC2::SecurityGroupIngress
    Properties:
      FromPort: 49152
      ToPort: 65535
      IpProtocol: tcp
      GroupId: !GetAtt FSxSecurityGroup.GroupId
      SourceSecurityGroupId: !GetAtt FSxSecurityGroup.GroupId
  TestFileSystem:
    Type: AWS::FSx::FileSystem
    Properties:
      FileSystemType: WINDOWS
      StorageCapacity: 32
      Tags:
        - Key: "Name"
          Value: "SampleFSxFileSystem"
      SecurityGroupIds:
        - !Ref FSxSecurityGroup
      SubnetIds:
        - !Ref SubnetId
      WindowsConfiguration:
        ActiveDirectoryId: !Ref ActiveDirectoryId
        DeploymentType: SINGLE_AZ_2
        ThroughputCapacity: 64
  RemotePowerShellExecutionInstanceRole:
    Type: AWS::IAM::Role
    Properties:
      Path: "/"
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: "Allow"
            Principal:
              Service:
                - "ec2.amazonaws.com"
                - "ssm.amazonaws.com"
            Action: "sts:AssumeRole"
      ManagedPolicyArns:
        - !Sub "arn:${AWS::Partition}:iam::aws:policy/AmazonSSMManagedInstanceCore"
        - !Sub "arn:${AWS::Partition}:iam::aws:policy/AmazonSSMDirectoryServiceAccess"
        - !Sub "arn:${AWS::Partition}:iam::aws:policy/CloudWatchAgentServerPolicy"
      Policies:
        - PolicyName: "AllowSsmExecution"
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: "Allow"
                Action: "ssm:CreateAssociation"
                Resource:
                  - !Sub "arn:${AWS::Partition}:ec2:${AWS::Region}:${AWS::AccountId}:instance/*"
        - PolicyName: "AllowGetParameter"
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: "Allow"
                Action: "ssm:GetParameter"
                Resource: !Sub "arn:${AWS::Partition}:ssm:${AWS::Region}:${AWS::AccountId}:parameter/${ActiveDirectoryCredentialsParameterName}"
  RemotePowerShellExecutionInstanceProfile:
    Type: AWS::IAM::InstanceProfile
    Properties:
      Path: "/"
      Roles:
        - !Ref RemotePowerShellExecutionInstanceRole
  RemotePowerShellExecutionInstance:
    Type: AWS::EC2::Instance
    Properties:
      ImageId: !Ref LatestWindowsAmiId
      IamInstanceProfile: !Ref RemotePowerShellExecutionInstanceProfile
      InstanceType: t3.small
      SecurityGroupIds: 
        - !Ref FSxSecurityGroup
      AvailabilityZone: !Ref ManagementInstanceAvailabilityZoneName
      SsmAssociations:
        - DocumentName: AWS-JoinDirectoryServiceDomain
          AssociationParameters:
            - Key: directoryId
              Value:
                - !Ref ActiveDirectoryId
            - Key: directoryName
              Value:
                - !Ref ActiveDirectoryName
            - Key: dnsIpAddresses
              Value: !Ref ActiveDirectoryDnsAddresses
  SsmExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      Path: "/"
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: "Allow"
            Principal:
              Service:
                - "ssm.amazonaws.com"
                - "ec2.amazonaws.com"
            Action: "sts:AssumeRole"
      ManagedPolicyArns:
        - !Sub "arn:${AWS::Partition}:iam::aws:policy/service-role/AmazonSSMAutomationRole"
  LambdaExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      Path: "/"
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: "Allow"
            Principal:
              Service: "lambda.amazonaws.com"
            Action: "sts:AssumeRole"
      ManagedPolicyArns:
        - !Sub "arn:${AWS::Partition}:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole"
      Policies:
        - PolicyName: "AllowSsmExecution"
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: "Allow"
                Action: "ssm:SendCommand"
                Resource: 
                  - !Sub "arn:${AWS::Partition}:ssm:*:${AWS::AccountId}:document/*"
                  - !Sub "arn:${AWS::Partition}:ec2:*:${AWS::AccountId}:instance/${RemotePowerShellExecutionInstance}"
        - PolicyName: "AllowDescribeFSxFileSystem"
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: "Allow"
                Action: "fsx:DescribeFileSystems"
                Resource: "*"
  FSxShadowCopies:
    Type: Custom::FSxShadowCopies
    Properties:
      ServiceToken: !GetAtt SsmInvocationLambdaFunction.Arn
      fsxFileSystemId: !Ref TestFileSystem
      storageCapacity: "16%"
      scheduleExpressions: 
        - "-weekly -DaysOfWeek Monday,Tuesday,Wednesday,Thursday,Friday -at 06:00"
        - "-weekly -DaysOfWeek Monday,Tuesday,Wednesday,Thursday,Friday -at 18:00"
      instanceId: !Ref RemotePowerShellExecutionInstance
      ssmDocName: !Ref RunPowerShellInvokeCommandDocument
      activeDirectoryCredentialsParameterName: !Ref ActiveDirectoryCredentialsParameterName
  SsmInvocationLambdaFunction:
    Type: AWS::Lambda::Function
    Properties:
      Handler: "index.handler"
      Role: !GetAtt LambdaExecutionRole.Arn
      Runtime: "python3.7"
      Timeout: 10
      Code:
        ZipFile: |
          import json
          import boto3
          import traceback
          import logging
          from botocore.vendored import requests

          logger = logging.getLogger()
          logger.setLevel(logging.INFO)
          ssmClient = boto3.client('ssm')

          def handler(event, context):

              logger.info(f'REQUEST BODY:\n{event}')

              fsxClient = boto3.client('fsx')
              fileSystemId = event["ResourceProperties"]["fsxFileSystemId"]
              fileSystem = fsxClient.describe_file_systems(FileSystemIds = [fileSystemId])
              fsxRemotePowershellEndpoint = fileSystem["FileSystems"][0]["WindowsConfiguration"]["RemoteAdministrationEndpoint"]
              instanceIds = [event["ResourceProperties"]["instanceId"]]
              activeDirectoryCredentialsParameterName = event["ResourceProperties"]["activeDirectoryCredentialsParameterName"]

              try:
                  eventType = event["RequestType"]

                  # if we get a delete from CloudFormation, set the FSx schedule and storage to 0 (which is what Default is)
                  if eventType == 'Delete':
                      removeShadowStorage = 'True'
                      scheduleExpressions = ''
                  # if we get a Create from CloudFormation, set the FSx schedule and storage
                  elif eventType == 'Create' or eventType == 'Update':
                      removeShadowStorage = 'False'
                      scheduleExpressions = '|'.join(event["ResourceProperties"].get("scheduleExpressions", []))

                  result = ssmClient.send_command(
                      DocumentName=event["ResourceProperties"]["ssmDocName"],
                      InstanceIds=instanceIds,
                      Parameters={"fsxRemoteEndpoint": [ fsxRemotePowershellEndpoint ], "removeShadowStorage": [removeShadowStorage], "storageCapacity": [event["ResourceProperties"]["storageCapacity"]], "scheduleExpressions": [scheduleExpressions], "activeDirectoryCredentialsParameterName": [activeDirectoryCredentialsParameterName]}
                  )

                  responseData = {'Result': result["Command"]["StatusDetails"], 'Version': ''}
                  responseStatus = 'SUCCESS'  
              except Exception as e:
                  logger.error(traceback.format_exc())
                  responseStatus = 'FAILED'
                  responseData = { 'Failure': f'Unable to fulfill {eventType} request. Check Lambda\'s CloudWatch log for details.' }

              sendResponse(event, context, responseStatus, responseData, fileSystemId)

          def sendResponse(event, context, responseStatus, responseData, physical_resource_id, reason=None):
              if 'StackId' not in event:
                  raise KeyError('Make sure that Lambda is being invoked with correct payload.')

              responseBody = {
                  'Status': responseStatus,
                  'Reason': 'See the details in CloudWatch Log Stream: ' + context.log_stream_name,
                  'PhysicalResourceId': physical_resource_id,
                  'StackId': event['StackId'],
                  'RequestId': event['RequestId'],
                  'LogicalResourceId': event['LogicalResourceId'],
                  'Data': responseData
              }

              logger.info('RESPONSE BODY:\n' + json.dumps(responseBody))
              responseUrl = event['ResponseURL']
              json_responseBody = json.dumps(responseBody)
              headers = { 'content-type': '', 'content-length': str(len(json_responseBody)) }

              try:
                  response = requests.put(responseUrl, data=json_responseBody, headers=headers)
                  logger.info(f'Status code: {response.reason}')
              except Exception as e:
                  logger.error('failed executing sendResponse: ' + traceback.format_exc())
  RunPowerShellInvokeCommandDocument:
    Type: AWS::SSM::Document
    Properties:
      DocumentType: Command
      Content:
        assumeRole: !GetAtt SsmExecutionRole.Arn
        description: Run a PowerShell Invoke-Command to start a psession and set the FSx Copy Schedule
        schemaVersion: "2.2"
        parameters:
          removeShadowStorage:
            type: "String"
          storageCapacity:
            type: "String"
          scheduleExpressions:
            type: "String"
          fsxRemoteEndpoint:
            type: String
            description: (Required) The Windows Remote PowerShell Endpoint for your FSx cluster
          activeDirectoryCredentialsParameterName:
            type: String
            description: (Required) The Systems Manager Parameter Store parameter name containing the Active Directory credentials JSON string
        mainSteps:
          - name: RunPowershellInvokeCommand
            action: "aws:runPowerShellScript"
            precondition:
              StringEquals:
                - "platformType"
                - "Windows"
            inputs:
              runCommand:
              - "try {"
              - "  $param = (Get-SSMParameter -Name '{{activeDirectoryCredentialsParameterName}}' -WithDecryption $True).Value | ConvertFrom-Json"
              - "  if (${{removeShadowStorage}}) {"
              - "    Invoke-Command -ComputerName {{fsxRemoteEndpoint}} -configurationname fsxremoteadmin -credential (New-Object System.Management.Automation.PSCredential $param.username,(ConvertTo-SecureString $param.password -AsPlainText -Force)) -ScriptBlock { Remove-FsxShadowStorage -Confirm:$False }"
              - "  }"
              - "  else {"
              - "    Invoke-Command -ComputerName {{fsxRemoteEndpoint}} -configurationname fsxremoteadmin -credential (New-Object System.Management.Automation.PSCredential $param.username,(ConvertTo-SecureString $param.password -AsPlainText -Force)) -ScriptBlock { Set-FsxShadowStorage -MaxSize '{{storageCapacity}}' }"
              - "    [CimInstance[]]$triggers = @()"
              - "    foreach ($triggerExpression in \"{{scheduleExpressions}}\".Split('|')) {"
              - "      $scriptBlockExpression = \"param([CimInstance[]]`$triggers) New-ScheduledTaskTrigger $triggerExpression\""
              - "      $scriptBlock = [ScriptBlock]::Create($scriptBlockExpression)"
              - "      $triggers += Invoke-Command -ScriptBlock $scriptBlock -ArgumentList $triggers"
              - "    }"
              - "    Invoke-Command -ComputerName {{fsxRemoteEndpoint}} -configurationname fsxremoteadmin -credential (New-Object System.Management.Automation.PSCredential $param.username,(ConvertTo-SecureString $param.password -AsPlainText -Force)) -ScriptBlock { Set-FsxShadowCopySchedule -scheduledtasktriggers $Using:triggers -Confirm:$False }"
              - "  }"
              - "}"
              - "catch {"
              - "  throw \"Error encountered when extracting sending commands to FSx file system.`n$($_.Exception.Message)\""
              - "}"
