AWSTemplateFormatVersion: "2010-09-09"
Description: Deploys CoreKMS server-signer Prime Onchain Wallet.

Parameters:
  SSHKeyName:
    Description: The EC2 Key Pair to allow SSH access to the instance
    Type: AWS::EC2::KeyPair::KeyName

  SSHIPRange:
    Description: The IP address range that can be used to SSH to your instance
    Type: String
    MinLength: 9
    MaxLength: 18
    Default: 10.0.0.0/8
    AllowedPattern: ^([0-9]{1,3}\.){3}[0-9]{1,3}(\/[0-9]+)?$
    ConstraintDescription: must be a valid IP CIDR range of the form x.x.x.x/x.

  ImageID:
    Description: The Amazon Machine Image ID to use for the instance. The default is the latest Amazon Linux 2 AMI in your region.
    Type: "AWS::SSM::Parameter::Value<AWS::EC2::Image::Id>"
    Default: "/aws/service/canonical/ubuntu/server/22.04/stable/current/amd64/hvm/ebs-gp2/ami-id"

Resources:
  ServerSignerEC2Instance:
    Type: AWS::EC2::Instance
    Properties:
      InstanceType: "t2.medium"
      KeyName: !Ref SSHKeyName
      ImageId: !Ref ImageID
      Tags:
        - Key: Name
          Value: corekms-server-signer
      UserData:
        Fn::Base64: !Sub |
          #!/bin/bash -xe
          # Wait Condition URL
          WAIT_CONDITION_URL="${WaitHandle}"

          function error_exit() {
            curl -X PUT -H 'Content-Type:' --data-binary '{"Status" : "FAILURE", "UniqueId": "SERVER_SIGNER_1234", "Reason" : "'"$1"'","Data" : "Application has completed configuration."}' \
                "$WAIT_CONDITION_URL"
            exit 1
          }

          sudo apt update && sudo apt-get install -y awscli golang-go postgresql-client jq unzip || error_exit "Failed to install packages"

          while ! pg_isready -d serversigner -h ${ServerSignerRDSCluster.Endpoint.Address} -p ${ServerSignerRDSCluster.Endpoint.Port}; do
            echo "Waiting for DB to respond..."
            sleep 10
          done

          TOKEN=$(curl --fail -X PUT "http://169.254.169.254/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 21600") || error_exit "Failed to get AWS token"
          region=$(curl --fail -H "X-aws-ec2-metadata-token: $TOKEN" "http://169.254.169.254/latest/meta-data/placement/region") || error_exit "Failed to get region"
          export region

          RDSMasterSecret=${AWS::StackName}-RDSMasterSecret
          masterpassword=$(aws secretsmanager get-secret-value --secret-id $RDSMasterSecret --region $region --query SecretString --output text 2>&1) || error_exit "Failed to get master password for $RDSMasterSecret: $masterpassword"
          if [[ -z "$masterpassword" ]]; then
            error_exit "Failed to get master password. password is empty"
          fi
          masterpassword=$(echo $masterpassword | jq -r '.password')
          # Wait until DB connection is available
          function check_db_connection() {
              PGPASSWORD=$(echo $masterpassword | tr -d '"') psql -h ${ServerSignerRDSCluster.Endpoint.Address} -p ${ServerSignerRDSCluster.Endpoint.Port} -U root -d serversigner -c '\q' &> /dev/null
          }

          retry_count=0
          MAX_RETRIES=10
          SLEEP_INTERVAL=5
          until check_db_connection; do
            retry_count=$((retry_count + 1))
            if [ "$retry_count" -ge "$MAX_RETRIES" ]; then
              error_exit "Failed to connect to database after $MAX_RETRIES attempts."
            fi
            echo "Waiting for db connection verify..."
            sleep "$SLEEP_INTERVAL"
          done
          # Configure DB connection
          PGPASSWORD=$(echo $masterpassword | tr -d '"') psql -h ${ServerSignerRDSCluster.Endpoint.Address} -p ${ServerSignerRDSCluster.Endpoint.Port} -U root -d serversigner -c "CREATE TABLE IF NOT EXISTS seeds (id SERIAL PRIMARY KEY,seed_id VARCHAR(255) NOT NULL,encrypted_seed TEXT NOT NULL, encrypted_kek TEXT NOT NULL);" || error_exit "Failed to create CDP tables"
          echo "Database connected successfully!"
          echo 'export SERVERSIGNER_DB_HOST="${ServerSignerRDSCluster.Endpoint.Address}"' >> /etc/profile || error_exit "Failed to set SERVERSIGNER_DB_HOST environment variables"
          echo 'export SERVERSIGNER_DB_PORT="${ServerSignerRDSCluster.Endpoint.Port}"' >> /etc/profile || error_exit "Failed to set SERVERSIGNER_DB_PORT environment variables"
          echo 'export SERVERSIGNER_DB_NAME="serversigner"' >> /etc/profile || error_exit "Failed to set SERVERSIGNER_DB_NAME environment variables"
          echo 'export SERVERSIGNER_DB_USER="root"' >> /etc/profile || error_exit "Failed to set SERVERSIGNER_DB_USER environment variables"
          # Password in secret manager could have special characters. So, we need to escape them.
          echo export SERVERSIGNER_DB_PASSWORD=$(echo $(printf '%q\n' $masterpassword)) >> /etc/profile || error_exit "Failed to set SERVERSIGNER_DB_PASSWORD environment variables"
          echo 'export SERVERSIGNER_KMS_KEY_ID="${ServerSignerEncKey.KeyId}"' >> /etc/profile || error_exit "Failed to set SERVERSIGNER_KMS_KEY_ID environment variables"
          echo export SERVERSIGNER_REGION="$region" >> /etc/profile || error_exit "Failed to set SERVERSIGNER_REGION environment variables"

          source /etc/profile || error_exit "Failed to source environment variables"

          # Configure KMS keys
          echo 'export SERVERSIGNER_SIGN_KMS_KEY_ID="${ServerSignerAuthKey.KeyId}"' >> /etc/profile || error_exit "Failed to set SERVERSIGNER_SIGN_KMS_KEY_ID environment variables"
          echo 'export SERVERSIGNER_SIGN_KMS_KEY_ARN="${ServerSignerAuthKey.Arn}"' >> /etc/profile || error_exit "Failed to set SERVERSIGNER_SIGN_KMS_KEY_ARN environment variables"
          kms_pub_key=$(aws kms get-public-key --key-id ${ServerSignerAuthKey.KeyId} --region $region | jq '.PublicKey' | tr -d '"') || error_exit "Failed to get KMS public key"
          echo export SERVERSIGNER_SIGN_KMS_KEY_PUBKEY="$kms_pub_key" >> /etc/profile || error_exit "Failed to set SERVERSIGNER_SIGN_KMS_KEY_PUBKEY environment variables"
          echo 'export SERVERSIGNER_ENCRYPT_KMS_KEY_ARN="${ServerSignerEncKey.Arn}"' >> /etc/profile || error_exit "Failed to set SERVERSIGNER_ENCRYPT_KMS_KEY_ARN environment variables"
          source /etc/profile || error_exit "Failed to source environment variables"

          # Download server-signer package
          wget -O cdp-signer.deb.zip https://api.cdp.coinbase.com/server-signer/assets/0.0.19/cdp-signer.deb.zip || error_exit "Failed to download signer"

          # Install server-signer
          sudo unzip -o cdp-signer.deb.zip && sudo rm cdp-signer.deb.zip || error_exit "Failed to unzip signer"
          sudo dpkg -i cdp-signer*.deb ||  error_exit "Failed to install signer"

          # Set up environment variables
          echo "APPROVER_DB_DATA_SOURCE=host=$SERVERSIGNER_DB_HOST port=$SERVERSIGNER_DB_PORT dbname=serversigner user=root password=$(echo $masterpassword | tr -d '"') sslmode=require" | sudo tee -a /etc/cdp-signer.env || error_exit "Failed to set APPROVER_DB_DATA_SOURCE environment variables"
          echo "APPROVER_KMS_SIGN_KEY_ID=$SERVERSIGNER_SIGN_KMS_KEY_ARN" | sudo tee -a /etc/cdp-signer.env || error_exit "Failed to set APPROVER_KMS_SIGN_KEY_ID environment variables"
          echo "APPROVER_KMS_SIGN_KEY_PUBKEY=$SERVERSIGNER_SIGN_KMS_KEY_PUBKEY" | sudo tee -a /etc/cdp-signer.env || error_exit "Failed to set APPROVER_KMS_SIGN_KEY_PUBKEY environment variables"
          echo "APPROVER_KMS_ENCRYPT_KEY_ID=$SERVERSIGNER_ENCRYPT_KMS_KEY_ARN" | sudo tee -a /etc/cdp-signer.env || error_exit "Failed to set APPROVER_KMS_ENCRYPT_KEY_ID environment variables"
          echo "APPROVER_KMS_ENCRYPT_KEY_ID=$SERVERSIGNER_ENCRYPT_KMS_KEY_ARN" | sudo tee -a /etc/cdp-signer.env || error_exit "Failed to set APPROVER_KMS_ENCRYPT_KEY_ID environment variables"
          # Disable refresh for speed
          echo "REFRESH_AFTER_DERIVE_SIGN=false" | sudo tee -a /etc/cdp-signer.env

          # Use CoreKMS 
          sudo sed -i 's/^APPROVER_ENV=.*/APPROVER_ENV=prod/' /etc/cdp-signer.env
          # Use the Prime partition
          sudo sed -i 's/^APPROVER_PARTITION=.*/APPROVER_PARTITION=INSTO_WEB3/' /etc/cdp-signer.env
          # Use two party MPC mode (seeds)
          sudo sed -i 's/^APPROVER_TWO_PARTY=.*/APPROVER_TWO_PARTY=true/' /etc/cdp-signer.env          
          # Set the AWS region
          echo AWS_REGION="$region" | sudo tee -a /etc/cdp-signer.env || error_exit "Failed to set AWS_REGION environment variables"

          # Provision
          source load-cdp-signer-env-vars || error_exit "Failed to load signer environment variables"
          enrollment_data=$(APPROVER_LOG_LEVEL=ERROR cdp-signer approver provision 2>&1) || error_exit "CoreKMS provision failed with $enrollment_data"
          echo $enrollment_data

          sudo systemctl start cdp-signer || error_exit "Failed to start server signer "

          # Send SUCCESS signal
          curl -X PUT -H 'Content-Type:' --data-binary "{\"Status\" : \"SUCCESS\",\"Reason\" : \"Server signer setup complete\", \"UniqueId\": \"$enrollment_data\" ,\"Data\" : \"\"}" "$WAIT_CONDITION_URL"

      SecurityGroups:
        - !Ref InstanceSecurityGroup
      IamInstanceProfile: !Ref InstanceProfile

  WaitHandle:
    Type: AWS::CloudFormation::WaitConditionHandle

  WaitCondition:
    Type: AWS::CloudFormation::WaitCondition
    DependsOn: ServerSignerEC2Instance
    Properties:
      Handle: !Ref WaitHandle
      Timeout: "10000" # Timeout in seconds

  InstanceProfile:
    Type: AWS::IAM::InstanceProfile
    Properties:
      Roles:
        - !Ref EC2AccessRole

  RDSMasterSecret:
    Type: AWS::SecretsManager::Secret
    Properties:
      Name: !Sub ${AWS::StackName}-RDSMasterSecret
      Description: "Master password for DB cluster"
      GenerateSecretString:
        SecretStringTemplate: '{"username":"root"}'
        GenerateStringKey: "password"
        PasswordLength: 16
        ExcludeCharacters: '"\\/@'

  ServerSignerRDSCluster:
    Type: AWS::RDS::DBCluster
    Properties:
      Engine: aurora-postgresql
      BackupRetentionPeriod: 1
      DatabaseName: serversigner
      DBClusterIdentifier: !Sub ${AWS::StackName}-serversigner-cluster
      MasterUsername: !Sub "{{resolve:secretsmanager:${RDSMasterSecret}:SecretString:username}}"
      MasterUserPassword: !Sub "{{resolve:secretsmanager:${RDSMasterSecret}:SecretString:password}}"
      VpcSecurityGroupIds:
        - !GetAtt ServerSignerDBSecurityGroup.GroupId

  ServerSignerRDSInstance:
    Type: AWS::RDS::DBInstance
    Properties:
      Engine: aurora-postgresql
      DBClusterIdentifier: !Ref ServerSignerRDSCluster
      DBInstanceClass: db.t3.medium

  ServerSignerEncKey:
    Type: AWS::KMS::Key
    Properties:
      KeyPolicy:
        Version: "2012-10-17"
        Id: !Sub ${AWS::StackName}-server-signer-enc-key
        Statement:
          - Sid: "Enable IAM User Permissions"
            Effect: Allow
            Principal:
              AWS: !Sub "arn:aws:iam::${AWS::AccountId}:root"
            Action: "kms:*"
            Resource: "*"
          - Sid: "Allow access for EC2 instances"
            Effect: Allow
            Principal:
              AWS: !GetAtt EC2AccessRole.Arn
            Action:
              - "kms:Encrypt"
              - "kms:Decrypt"
            Resource: "*"
      Description: !Sub ${AWS::StackName}-Server signer key for encrypting key-material data with envelope encryption (KEK)
      KeyUsage: "ENCRYPT_DECRYPT"
      Origin: "AWS_KMS"
      Enabled: true

  ServerSignerAuthKey:
    Type: AWS::KMS::Key
    Properties:
      KeyPolicy:
        Version: "2012-10-17"
        Id: !Sub ${AWS::StackName}-server-signer-auth-key
        Statement:
          - Sid: "Enable IAM User Permissions"
            Effect: Allow
            Principal:
              AWS: !Sub "arn:aws:iam::${AWS::AccountId}:root"
            Action: "kms:*"
            Resource: "*"
          - Sid: "Allow access for EC2 instances"
            Effect: Allow
            Principal:
              AWS: !GetAtt EC2AccessRole.Arn
            Action:
              - "kms:Sign"
              - "kms:Verify"
              - "kms:DescribeKey"
              - "kms:GetPublicKey"
            Resource: "*"
      Description: !Sub ${AWS::StackName}-Server signer key for secure-channel authentication
      KeySpec: "ECC_NIST_P256"
      KeyUsage: "SIGN_VERIFY"
      Origin: "AWS_KMS"
      Enabled: true

  EC2AccessRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service: [ec2.amazonaws.com]
            Action: [sts:AssumeRole]
      Policies:
        - PolicyName: !Sub ${AWS::StackName}-SecretAccessPolicy
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action:
                  - "secretsmanager:GetSecretValue"
                Resource: !Sub "arn:aws:secretsmanager:${AWS::Region}:${AWS::AccountId}:secret:${AWS::StackName}-RDSMasterSecret-*"
        - PolicyName: !Sub ${AWS::StackName}-CloudWatchAgentPolicy
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action:
                  - "logs:CreateLogGroup"
                  - "logs:CreateLogStream"
                  - "logs:PutLogEvents"
                Resource: "*"

  InstanceSecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupDescription: Enable SSH access via port 22
      SecurityGroupIngress:
        - IpProtocol: tcp
          FromPort: 22
          ToPort: 22
          CidrIp: !Ref SSHIPRange

  ServerSignerDBSecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupDescription: "Security group for Aurora Cluster"
      SecurityGroupIngress:
        - IpProtocol: tcp
          FromPort: 5432
          ToPort: 5432
          SourceSecurityGroupId: !GetAtt InstanceSecurityGroup.GroupId

Outputs:
  EnrollData:
    Description: "Participant enrollment data"
    Value: !GetAtt WaitCondition.Data