Description: Research and Engineering Studio SSO setup with Keycloak

Metadata:
  AWS::CloudFormation::Interface:
    ParameterGroups:
      - Label:
          default: RES SSO Configuration
        Parameters:
          - EnvironmentName
          - Keypair
          - ServiceAccountCredentialsSecretArn
          - VpcId
          - PublicSubnet
          - ServiceAccountUserDN
          - UsersDN
          - LDAPConnectionURI

Parameters:

  EnvironmentName:
    Description: Provide name of the RES Environment. Must be unique for your account and AWS Region.
    Type: String
    Default: res-demo
    AllowedPattern: ^res-[a-z\-\_0-9]{0,7}$
    ConstraintDescription: EnvironmentName must start with "res-" without capital letters and should be less than or equal to 11 characters.

  Keypair:
    Description: EC2 Keypair to access management instance.
    Type: AWS::EC2::KeyPair::KeyName

  ServiceAccountCredentialsSecretArn:
    Type: String
    AllowedPattern: ^(?:arn:(?:aws|aws-us-gov|aws-cn):secretsmanager:[a-z0-9-]{1,20}:[0-9]{12}:secret:[A-Za-z0-9-_+=,\.@]{1,128})?$
    Description: Directory Service Service Account Credentials Secret ARN. The username and password for the Active Directory ServiceAccount user formatted as a username:password key/value pair.

  VpcId:
    Type: AWS::EC2::VPC::Id
    AllowedPattern: vpc-[0-9a-f]{17}
    ConstraintDescription: VpcId must begin with 'vpc-', only contain letters (a-f) or numbers(0-9) and must be 21 characters in length
  
  PublicSubnet:
    Type: AWS::EC2::Subnet::Id
    AllowedPattern: subnet-.+
    Description: Select a public subnet from the already selected VPC
  
  ServiceAccountUserDN:
    Type: String
    AllowedPattern: .+
    Description: Provide the Distinguished name (DN) of the service account user in the Active Directory
  
  UsersDN:
    Type: String
    AllowedPattern: .+
    Description: Please provide Users Organization Unit in your active directory under which all of your users exist. For example, OU=Users,DC=RES,DC=example,DC=internal

  LDAPConnectionURI:
    Type: String
    AllowedPattern: .+
    Description: Please provide the active directory connection URI (e.g. ldap://www.example.com)

Resources:

  Keycloak:
    Type: AWS::CloudFormation::Stack
    Properties:
      Parameters:
        Keypair: !Ref Keypair
        ServiceAccountCredentialsSecretArn: !Ref ServiceAccountCredentialsSecretArn
        VpcId: !Ref VpcId
        PublicSubnet: !Ref PublicSubnet
        ServiceAccountUserDN: !Ref ServiceAccountUserDN
        UsersDN: !Ref UsersDN
        LDAPConnectionURI: !Ref LDAPConnectionURI
        CogntioUserPoolId: !Sub ${DataGatherCustomResource.UserPoolId}
        EnvironmentBaseURL: !Sub ${DataGatherCustomResource.LoadBalancerDnsName}
        SAMLRedirectUrl: !Sub ${DataGatherCustomResource.SAMLRedirectUrl}
      TemplateURL: https://aws-hpc-recipes.s3.us-east-1.amazonaws.com/main/recipes/res/res_demo_env/assets/keycloak.yaml

  KeycloakDataGatherLambdaExecutionRole:
    Type: 'AWS::IAM::Role'
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: 'sts:AssumeRole'
      Policies:
        - PolicyName: QueryCognitoAndELBv2
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - elasticloadbalancing:DescribeLoadBalancers
                  - elasticloadbalancing:DescribeTags
                  - cognito-idp:ListUserPools
                Resource: '*'
              - Effect: Allow
                Action:
                  - cognito-idp:DescribeUserPool
                Resource:
                  - !Sub arn:${AWS::Partition}:cognito-idp:${AWS::Region}:${AWS::AccountId}:userpool/${AWS::Region}*
              - Effect: Allow
                Action:
                  - logs:CreateLogGroup
                  - logs:CreateLogStream
                  - logs:PutLogEvents
                Resource: '*'

  KeycloakDataGatherHandlerFunction:
    Type: 'AWS::Lambda::Function'
    DependsOn:
     - KeycloakDataGatherLambdaExecutionRole
    Properties:
      Description: 'Keycloak Data Gather Handler'
      FunctionName: !Sub KeycloakDataGatherHandler-${EnvironmentName}
      Timeout: 300 # 5 minutes
      Role: !GetAtt KeycloakDataGatherLambdaExecutionRole.Arn
      Handler: index.handler
      Runtime: python3.11
      Code:
        ZipFile: |
          import os
          import boto3
          import urllib.error
          import urllib.parse
          import urllib.request
          import json
          from typing import Any, Dict, TypedDict, Union
          from itertools import chain
          import boto3
          import botocore.exceptions
          import logging
          from typing import TypedDict
          
          logger = logging.getLogger()
          logger.setLevel(logging.INFO)

          TAG_NAME = "res:EnvironmentName"
          
          class CustomResourceResponse(TypedDict):
            Status: str
            Reason: str
            PhysicalResourceId: str
            StackId: str
            RequestId: str
            LogicalResourceId: str
          
          def send_response(url, response):
            request = urllib.request.Request(
              method="PUT",
              url=url,
              data=json.dumps(response).encode("utf-8"),
            )
            urllib.request.urlopen(request)
          
          def get_cognito_data(cluster_name, region_name):
            cognito_client = boto3.client("cognito-idp")
            logger.info(f"Working on getting Cognito details")
            userpool_pagintor = cognito_client.get_paginator("list_user_pools")
            userpool_pages = map(lambda p: p.get("UserPools", []), userpool_pagintor.paginate(MaxResults=50))
            userpool_match_fn = lambda up: up.get("Name", "") == f"{cluster_name}-user-pool" and up.get("Id")
            userpools = filter(userpool_match_fn, chain.from_iterable(userpool_pages))
            for userpool in userpools:
              pool_name = userpool.get("Name")
              pool_id = userpool.get("Id")
              logger.info(f"Processing cognito pool: {pool_name} PoolId: {pool_id}")
              describe_user_pool_result = cognito_client.describe_user_pool(
                UserPoolId=pool_id
              )
              tags = describe_user_pool_result.get("UserPool", {}).get(
                "UserPoolTags", {}
              )
              match_fn = lambda tg: tg[0] == TAG_NAME and tg[1] == cluster_name
              if next(filter(match_fn, tags.items()), None):
                logger.info(f"Found matching tags")
                domain = describe_user_pool_result['UserPool']['Domain']
                saml_redirect_url = f'https://{domain}.auth.{region_name}.amazoncognito.com/saml2/idpresponse'
                return pool_id, saml_redirect_url
              else:
                logger.info("No matching tags found")

          def get_alb_dns(cluster_name):
            elbv2_client = boto3.client('elbv2')
            logger.info(f"Working on getting load balancer DNS")
            lb_paginator = elbv2_client.get_paginator("describe_load_balancers")
            lb_pages = map(lambda p: p.get("LoadBalancers", []), lb_paginator.paginate())
            lb_match_fn = lambda lb: lb.get("LoadBalancerName", "") == f"{cluster_name}-external-alb"
            load_balancers = filter(lb_match_fn, chain.from_iterable(lb_pages))
            for load_balancer in load_balancers:
              load_balancer_arn = load_balancer.get("LoadBalancerArn", "")
              load_balancer_name = load_balancer.get("LoadBalancerName", "")
              logger.info(f"Processing load balancer: {load_balancer_name}")
              tag_description = elbv2_client.describe_tags(ResourceArns=[load_balancer_arn]).get('TagDescriptions', [None])[0]
              tags = tag_description.get('Tags', [])
              match_fn = lambda t: t['Key'] == TAG_NAME and t['Value'] == cluster_name
              if next(filter(match_fn, tags), None):
                logger.info(f"Found matching tags")
                return f'https://{load_balancer["DNSName"]}'
              else:
                logger.info("No matching tags found")
          
          def handler(event, _):
            logger.info(f"Received event: {event}")
            request_type = event["RequestType"]
            response_url = event["ResponseURL"]
            response = CustomResourceResponse(
              Status="SUCCESS",
              Reason="SUCCESS",
              PhysicalResourceId=event["LogicalResourceId"],
              StackId=event["StackId"],
              RequestId=event["RequestId"],
              LogicalResourceId=event["LogicalResourceId"],
              Data={}
            )
            if request_type == "Delete":
              send_response(response_url, response)
              return
            
            cluster_name = os.environ['CLUSTER_NAME']
            region_name = os.environ['AWS_REGION']

            try:
              user_pool_id, saml_redirect_url = get_cognito_data(cluster_name, region_name)
              dns_name = get_alb_dns(cluster_name)
              if not user_pool_id or not saml_redirect_url or not dns_name:
                raise Exception(f"Unable to find matching cognito user pool, SAML redirect URL for the user pool, or load balancer. Response: {response}")
              response["Data"]["UserPoolId"] = user_pool_id
              response["Data"]["SAMLRedirectUrl"] = saml_redirect_url
              response["Data"]["LoadBalancerDnsName"] = dns_name
            except Exception as e:
              logger.error(f"Error processing request {e}")
              response["Status"] = "FAILED"
              response["Reason"] = str(e)
            finally:
              logger.info(f"Sending response: {response}")
              send_response(url=response_url, response=response)
      Environment:
        Variables:
          CLUSTER_NAME: !Ref EnvironmentName

  DataGatherCustomResource:
    Type: Custom::KeycloakDataGather
    Properties:
      ServiceToken: !GetAtt KeycloakDataGatherHandlerFunction.Arn
  
  InvokeConfigureSSOLambdaRole:
    Type: 'AWS::IAM::Role'
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: 'sts:AssumeRole'
      Policies:
        - PolicyName: InvokeConfigureSSOLambdaPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - lambda:InvokeFunction
                Resource:
                  - !Sub arn:${AWS::Partition}:lambda:${AWS::Region}:${AWS::AccountId}:function:${EnvironmentName}-configure_sso
              - Effect: Allow
                Action:
                  - logs:CreateLogGroup
                  - logs:CreateLogStream
                  - logs:PutLogEvents
                Resource: '*'

  InvokeConfigureSSOHandlerFunction:
    Type: 'AWS::Lambda::Function'
    DependsOn: 
     - InvokeConfigureSSOLambdaRole
     - Keycloak
    Properties:
      Description: 'Invoke RES configure sso function'
      FunctionName: !Sub InvokeConfigureSSOHandlerFunction-${EnvironmentName}
      Timeout: 300 # 5 minutes
      Role: !GetAtt InvokeConfigureSSOLambdaRole.Arn
      Handler: index.handler
      Runtime: python3.11
      Code:
        ZipFile: |
          import os
          import boto3
          import urllib.error
          import urllib.parse
          import urllib.request
          import json
          from typing import Any, Dict, TypedDict, Union

          import boto3
          import botocore.exceptions
          import base64
          import logging
          from typing import TypedDict
          
          logger = logging.getLogger()
          logger.setLevel(logging.INFO)
          
          class CustomResourceResponse(TypedDict):
            Status: str
            Reason: str
            PhysicalResourceId: str
            StackId: str
            RequestId: str
            LogicalResourceId: str
          
          def send_response(url, response):
            request = urllib.request.Request(
              method="PUT",
              url=url,
              data=json.dumps(response).encode("utf-8"),
            )
            urllib.request.urlopen(request)
          
          def handler(event, _):
            logger.info(f"Received event: {event}")

            try:
              request_type = event["RequestType"]
              response_url = event["ResponseURL"]
              response = CustomResourceResponse(
                Status="SUCCESS",
                Reason="SUCCESS",
                PhysicalResourceId=event["LogicalResourceId"],
                StackId=event["StackId"],
                RequestId=event["RequestId"],
                LogicalResourceId=event["LogicalResourceId"],
                Data={}
              )
              if request_type == "Delete":
                send_response(response_url, response)
                return
              lambda_name = os.environ['LAMBDA_NAME']
              region_name = os.environ['AWS_REGION']
              keycloak_url = os.environ['KEYCLOAK_URL']

              #Get SAML metadata string from Keycloak
              saml_metadata_url = f"{keycloak_url}/realms/res/protocol/saml/descriptor"
              logger.info(f"SAML metadata url: {saml_metadata_url}")

              local_filename, headers = urllib.request.urlretrieve(saml_metadata_url)
              saml_metadata = open(local_filename, "r").read()
              saml_metadata_utf8encoded = saml_metadata.encode("utf-8")
              saml_metadata_base64_bytes = base64.b64encode(saml_metadata_utf8encoded)
              saml_metadata_base64_string = saml_metadata_base64_bytes.decode("utf-8")

              #Build payload
              payload = json.dumps({ 
                'configure_sso_request': {
                  'provider_name': 'idc', 
                  'provider_type': 'SAML', 
                  'provider_email_attribute': 'email', 
                  'saml_metadata_file': saml_metadata_base64_string
                }
              })

              #Invoke Lambda
              logger.info(f"Invoking configure_sso lambda with payload : {payload}")
              lambda_client = boto3.client("lambda")
              lambda_response = lambda_client.invoke(
                FunctionName=lambda_name,
                Payload=payload
              )

              logger.info(f"Response from configure_sso lambda: lambda_response")
              if 'FunctionError' in lambda_response:
                response_payload = json.loads(response['Payload'].read())
                if 'errorMessage' in response_payload:
                  raise Exception(response_payload['errorMessage'])
                raise Exception(lambda_response['FunctionError'])
            except Exception as e:
              logger.error(f"Error processing request {e}")
              response["Status"] = "FAILED"
              response["Reason"] = str(e)
            finally:
              logger.info(f"Sending response: {response}")
              send_response(url=response_url, response=response)
      Environment:
        Variables:
          LAMBDA_NAME: !Sub ${EnvironmentName}-configure_sso
          KEYCLOAK_URL: !GetAtt [ Keycloak, Outputs.KeycloakUrl ]

  InvokeConfigureSSOCustomResource:
    Type: Custom::InvokeConfigureSSO
    Properties:
      ServiceToken: !GetAtt InvokeConfigureSSOHandlerFunction.Arn

Outputs:
  KeycloakUrl:
    Description: Keycloak Administrator Url
    Value: !GetAtt [ Keycloak, Outputs.KeycloakUrl ]
  KeycloakAdminPasswordSecretArn:
    Description: Keycloak password for admin user
    Value: !GetAtt [ Keycloak, Outputs.KeycloakAdminPasswordSecretArn ]
  ApplicationUrl:
    Description: RES application Url
    Value: !GetAtt DataGatherCustomResource.LoadBalancerDnsName