# source: https://aws-hpc-recipes.s3.amazonaws.com/main/recipes/pcs/getting_started/assets/pcs-cluster-sg.yaml

AWSTemplateFormatVersion: 2010-09-09
Description: Security group for AWS PCS clusters.
  This template creates a self-referencing security group that enables communications between AWS PCS controller, compute nodes, and client nodes.
  Optionally, it can also create a security group to enable SSH access to the cluster, and DCV remote desktop access to the login node.
  Check the Outputs tab of this stack for useful details about resources created by this template.

Metadata:
  AWS::CloudFormation::Interface:
    ParameterGroups:
      - Label:
          default: Network
        Parameters:
          - VpcId
      - Label:
          default: Security group configuration
        Parameters:
          - CreateInboundSshSecurityGroup
          - CreateInboundDcvSecurityGroup
          - ClientIpCidr

Parameters:
  VpcId:
    Description: VPC where the AWS PCS cluster will be deployed
    Type: 'AWS::EC2::VPC::Id'
  ClientIpCidr:
    Description: IP address(s) allowed to connect to nodes using SSH or DCV.
    Default: '0.0.0.0/0'
    Type: String
    AllowedPattern: (\d{1,3})\.(\d{1,3})\.(\d{1,3})\.(\d{1,3})/(\d{1,2})
    ConstraintDescription: Value must be a valid IP or network range of the form x.x.x.x/x.
  CreateInboundSshSecurityGroup:
    Description: Create an inbound security group to allow SSH access to nodes.
    Type: String
    Default: 'True'
    AllowedValues:
      - 'True'
      - 'False'
  CreateInboundDcvSecurityGroup:
    Description: Create an inbound security group to allow DCV access to login nodes on TCP/UDP 8443.
    Type: String
    Default: 'False'
    AllowedValues:
      - 'True'
      - 'False'

Conditions:
  CreateSshSecGroup: !Equals [!Ref CreateInboundSshSecurityGroup, 'True']
  CreateDcvSecGroup: !Equals [!Ref CreateInboundDcvSecurityGroup, 'True']

Resources:

  ClusterSecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupDescription: Supports communications between AWS PCS controller, compute nodes, and client nodes
      VpcId: !Ref VpcId
      GroupName: !Sub 'cluster-${AWS::StackName}'

  ClusterAllowAllInboundFromSelf:
    Type: AWS::EC2::SecurityGroupIngress
    Properties:
      GroupId: !Ref ClusterSecurityGroup
      IpProtocol: '-1'
      SourceSecurityGroupId: !Ref ClusterSecurityGroup

  ClusterAllowAllOutboundToSelf:
    Type: AWS::EC2::SecurityGroupEgress
    Properties:
      GroupId: !Ref ClusterSecurityGroup
      IpProtocol:  '-1'
      DestinationSecurityGroupId: !Ref ClusterSecurityGroup

  # This allows all outbound comms, which enables HTTPS calls and connections to networked storage
  ClusterAllowAllOutboundToWorld:
    Type: AWS::EC2::SecurityGroupEgress
    Properties:
      GroupId: !Ref ClusterSecurityGroup
      IpProtocol: '-1'
      CidrIp: 0.0.0.0/0

  # Attach this to login nodes to enable inbound SSH access.
  InboundSshSecurityGroup:
    Condition: CreateSshSecGroup
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupDescription: Allows inbound SSH access
      GroupName: !Sub 'inbound-ssh-${AWS::StackName}'
      VpcId: !Ref VpcId
      SecurityGroupIngress:
        - IpProtocol: tcp
          FromPort: 22
          ToPort: 22
          CidrIp: !Ref ClientIpCidr

  # Attach this to login nodes to enable inbound DCV access on TCP/UDP 8443.
  InboundDcvSecurityGroup:
    Condition: CreateDcvSecGroup
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupDescription: Allows inbound DCV access on TCP/UDP 8443
      GroupName: !Sub 'inbound-dcv-${AWS::StackName}'
      VpcId: !Ref VpcId
      SecurityGroupIngress:
        - IpProtocol: tcp
          FromPort: 8443
          ToPort: 8443
          CidrIp: !Ref ClientIpCidr
        - IpProtocol: udp
          FromPort: 8443
          ToPort: 8443
          CidrIp: !Ref ClientIpCidr

Outputs:
  ClusterSecurityGroupId:
    Description: Supports communication between PCS controller, compute nodes, and login nodes
    Value: !Ref ClusterSecurityGroup
  InboundSshSecurityGroupId:
    Condition: CreateSshSecGroup
    Description: Enables SSH access to login nodes
    Value: !Ref InboundSshSecurityGroup
  InboundDcvSecurityGroupId:
    Condition: CreateDcvSecGroup
    Description: Enables DCV access to login nodes on TCP/UDP 8443
    Value: !Ref InboundDcvSecurityGroup