AWSTemplateFormatVersion: 2010-09-09
Description: Kali Linux with NICE DCV (login as kali)

Metadata:
  License:
    Description: >
      Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
      SPDX-License-Identifier: MIT-0

      Permission is hereby granted, free of charge, to any person obtaining a copy of this
      software and associated documentation files (the "Software"), to deal in the Software
      without restriction, including without limitation the rights to use, copy, modify,
      merge, publish, distribute, sublicense, and/or sell copies of the Software, and to
      permit persons to whom the Software is furnished to do so.

      THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
      INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
      PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
      HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
      OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
      SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
  AWS::CloudFormation::Interface:
    ParameterGroups:
      - Label:
          default: AMI and instance type. Subscribe before launching ( https://aws.amazon.com/marketplace/pp/prodview-fznsw3f7mq7to )
        Parameters:
          - imageId
          - instanceType
      - Label:
          default: EC2 configuration
        Parameters:
          - ec2Name
          - vpcID
          - subnetID
          - assignStaticIP
          - displayPublicIP
      - Label:
          default: Allowed source IP prefixes
        Parameters:
          - ingressIPv4
          - ingressIPv6
      - Label:
          default: EBS volume configuration
        Parameters:
          - volumeSize
          - deviceName
      - Label:
          default: NICE DCV configuration
        Parameters:
          - listenPort

Parameters:
  imageId:
    Type: AWS::SSM::Parameter::Value<AWS::EC2::Image::Id>
    Description: Kali AMI ID ( aws ssm get-parameters-by-path --path /aws/service/marketplace/prod-tsqyof4l3a3aa/ --query "Parameters[].Name" )
    Default: /aws/service/marketplace/prod-tsqyof4l3a3aa/latest
  instanceType:
    Type: String
    Description: Instance type ( https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instance-types.html )
    Default: t3.medium

  ec2Name:
    Type: String
    Description: EC2 instance name
    Default: Kali Linux-NICE-DCV
  vpcID:
    Type: AWS::EC2::VPC::Id
    Description: VPC with internet connectivity
    ConstraintDescription: Do specify a valid value
    AllowedPattern: ".+"
  subnetID:
    Type: AWS::EC2::Subnet::Id
    Description: Subnet with internet connectivity
    ConstraintDescription: Do specify a valid value
    AllowedPattern: ".+"
  assignStaticIP:
    Type: String
    Description: Associate static public IPv4 address ( https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/elastic-ip-addresses-eip.html )
    AllowedValues:
      - "Yes"
      - "No"
    Default: "No"
  displayPublicIP:
    Type: String
    Description: Display EC2 public IP in CloudFormation Outputs (select No if EC2 has no public IP)
    AllowedValues:
      - "Yes"
      - "No"
    Default: "Yes"

  ingressIPv4:
    Type: String
    Description: Allowed source prefix (IPv4) ( e.g. 1.2.3.4/32, get your source IP from https://checkip.amazonaws.com )
    Default: 0.0.0.0/0
  ingressIPv6:
    Type: String
    Description: Allowed source prefix (IPv6)
    Default: ::/0

  listenPort:
    Type: Number
    Description: NICE DCV server TCP/UDP port
    ConstraintDescription: Specify a value between 1024 and 65535
    MinValue: 1024
    MaxValue: 65535
    Default: 8443

  volumeSize:
    Type: Number
    Description: Volume Size in GiBs (must be equal or larger than snapshot size)
    Default: 15
  deviceName:
    Type: String
    Description: Device Name
    Default: /dev/xvda

Conditions:
  useElasticIP: !Equals [!Ref assignStaticIP, "Yes"]
  displayPublicIP: !Equals [!Ref displayPublicIP, "Yes"]

Resources:
  securityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupDescription: Allow inbound DCV
      VpcId: !Ref vpcID
      SecurityGroupIngress:
        - Description: NICE DCV (IPv4)
          IpProtocol: "tcp"
          FromPort: !Ref listenPort
          ToPort: !Ref listenPort
          CidrIp: !Ref ingressIPv4
        - Description: NICE DCV QUIC (IPv4)
          IpProtocol: "udp"
          FromPort: !Ref listenPort
          ToPort: !Ref listenPort
          CidrIp: !Ref ingressIPv4
        - Description: NICE DCV (IPv6)
          IpProtocol: "tcp"
          FromPort: !Ref listenPort
          ToPort: !Ref listenPort
          CidrIpv6: !Ref ingressIPv6
        - Description: NICE DCV QUIC (IPv6)
          IpProtocol: "udp"
          FromPort: !Ref listenPort
          ToPort: !Ref listenPort
          CidrIpv6: !Ref ingressIPv6
      SecurityGroupEgress:
        - Description: Allow all outbound traffic (IPv4)
          IpProtocol: "-1"
          CidrIp: 0.0.0.0/0
        - Description: Allow all outbound traffic (IPv6)
          IpProtocol: "-1"
          CidrIpv6: "::/0"
      Tags:
        - Key: StackName
          Value: !Sub ${AWS::StackName}
        - Key: StackId
          Value: !Sub ${AWS::StackId}
        - Key: Name
          Value: !Sub "[${AWS::StackName}] - ${ec2Name}"
        - Key: GitHub
          Value: https://github.com/aws-samples/amazon-ec2-nice-dcv-samples

  instanceIamRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Statement:
          - Effect: Allow
            Principal:
              Service: [ec2.amazonaws.com]
            Action: ["sts:AssumeRole"]
      Path: /
      Policies: # https://docs.aws.amazon.com/dcv/latest/adminguide/setting-up-license.html
        - PolicyName: dcvLicensing
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action:
                  - s3:GetObject
                Resource: !Sub "arn:${AWS::Partition}:s3:::dcv-license.${AWS::Region}/*"
        - PolicyName: gpuDrivers
          PolicyDocument: # https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/install-nvidia-driver.html https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/install-amd-driver.html
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action:
                  - s3:Get*
                  - s3:List*
                Resource:
                  - !Sub "arn:${AWS::Partition}:s3:::nvidia-gaming"
                  - !Sub "arn:${AWS::Partition}:s3:::nvidia-gaming/*"
                  - !Sub "arn:${AWS::Partition}:s3:::ec2-linux-nvidia-drivers"
                  - !Sub "arn:${AWS::Partition}:s3:::ec2-linux-nvidia-drivers/*"
                  - !Sub "arn:${AWS::Partition}:s3:::ec2-amd-linux-drivers"
                  - !Sub "arn:${AWS::Partition}:s3:::ec2-amd-linux-drivers/*"
      ManagedPolicyArns:
        - !Sub "arn:${AWS::Partition}:iam::aws:policy/AmazonSSMManagedInstanceCore"
      Tags:
        - Key: StackName
          Value: !Sub ${AWS::StackName}
        - Key: StackId
          Value: !Sub ${AWS::StackId}
        - Key: GitHub
          Value: https://github.com/aws-samples/amazon-ec2-nice-dcv-samples

  instanceProfile:
    Type: AWS::IAM::InstanceProfile
    Properties:
      Path: /
      Roles:
        - !Ref instanceIamRole

  ec2Instance:
    Type: AWS::EC2::Instance
    CreationPolicy:
      ResourceSignal:
        Timeout: PT90M
    Properties:
      ImageId: !Ref imageId
      InstanceType: !Ref instanceType
      IamInstanceProfile: !Ref instanceProfile
      SubnetId: !Ref subnetID
      Monitoring: true
      SecurityGroupIds:
        - !Ref securityGroup
      BlockDeviceMappings:
        - DeviceName: !Ref deviceName
          Ebs:
            VolumeType: gp3
            VolumeSize: !Ref volumeSize
            DeleteOnTermination: true
      UserData:
        Fn::Base64: !Sub |
          #!/bin/zsh
          cd /root/

          export DEBIAN_FRONTEND=noninteractive
          apt-get update

          # https://docs.aws.amazon.com/systems-manager/latest/userguide/agent-install-ubuntu.html
          apt-get install -q -y snapd
          systemctl enable snapd
          systemctl start snapd
          sleep 2
          snap install amazon-ssm-agent --classic
          snap start amazon-ssm-agent

          apt-get install -q -y wget tmux unzip tar curl sed
          # https://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/releasehistory-aws-cfn-bootstrap.html#releasehistory-aws-cfn-bootstrap-v1
          # https://aws.amazon.com/premiumsupport/knowledge-center/install-cloudformation-scripts/        
          apt-get install -q -y python3 python3-setuptools python3-docutils python3-daemon
          wget -4 -nv https://s3.amazonaws.com/cloudformation-examples/aws-cfn-bootstrap-py3-latest.tar.gz
          tar -xvf aws-cfn-bootstrap-py3-latest.tar.gz
          cd aws-cfn-bootstrap-2.0
          python3 setup.py build
          python3 setup.py install

          cd /root/
          # cfn-init
          python3 /usr/local/bin/cfn-init -v --stack ${AWS::StackName} --resource ec2Instance --region ${AWS::Region}

          apt-get -q -y install kali-desktop-xfce
          apt-get -q -y install pulseaudio-utils

          # https://docs.aws.amazon.com/dcv/latest/adminguide/setting-up-installing-linux-server.html
          wget -nv https://d1uj6qtbmh3dt5.cloudfront.net/NICE-GPG-KEY
          gpg --import NICE-GPG-KEY
          wget -nv https://d1uj6qtbmh3dt5.cloudfront.net/nice-dcv-ubuntu2204-x86_64.tgz
          tar -xvzf nice-dcv-ubuntu2204-x86_64.tgz && cd nice-dcv-*-ubuntu2204-x86_64

          # tweaks for installation on Kali
          ln -s /etc/os-release /etc/lsb-release
          mkdir -p /etc/lightdm/lightdm.conf.d

          # https://docs.aws.amazon.com/dcv/latest/adminguide/setting-up-installing-linux-server.html#linux-server-install
          apt-get -q -y install ./nice-dcv-server_*_amd64.ubuntu2204.deb
          apt-get -q -y install ./nice-dcv-web-viewer_*_amd64.ubuntu2204.deb
          usermod -aG video dcv 
          # from /etc/lightdm/lightdm.conf.d
          sed -i '/^\[Seat\:\*\]/a display-setup-script=/usr/lib/x86_64-linux-gnu/dcv/dcvlightdm' /etc/lightdm/lightdm.conf

          # virtual session support
          apt-get -q -y install ./nice-xdcv_*_amd64.ubuntu2204.deb

          # https://docs.aws.amazon.com/dcv/latest/adminguide/enable-quic.html
          cp /etc/dcv/dcv.conf /etc/dcv/dcv.conf.org
          sed -i "s/^#enable-quic-frontend=true/enable-quic-frontend=true/g" /etc/dcv/dcv.conf

          # session storage: https://docs.aws.amazon.com/dcv/latest/userguide/using-transfer.html
          # https://docs.aws.amazon.com/dcv/latest/adminguide/managing-sessions-start.html#managing-sessions-start-manual
          cat << EoF > /etc/systemd/system/dcv-virtual-session.service
          [Unit]
          Description=Create DCV virtual session
          After=default.target network.target 

          [Service]
          ExecStart=/opt/dcv-virtual-session.sh 

          [Install]
          WantedBy=default.target
          EoF

          cat << EoF > /opt/dcv-virtual-session.sh
          #!/bin/zsh
          dcvUser=kali
          while true;
          do
            if (/usr/bin/dcv list-sessions | grep \$dcvUser 1>/dev/null)
            then
              sleep 5
            else
              /usr/bin/dcv create-session \$dcvUser --owner \$dcvUser --storage-root /home/\$dcvUser
              /usr/bin/dcv list-sessions
            fi
          done
          EoF
          chmod +x /opt/dcv-virtual-session.sh

          # https://docs.aws.amazon.com/dcv/latest/adminguide/manage-port-addr.html
          sed -i "/^web-port=/d" /etc/dcv/dcv.conf
          sed -i "/^quic-port=/d" /etc/dcv/dcv.conf
          sed -i "/^\[connectivity\]/a web-port=${listenPort}" /etc/dcv/dcv.conf
          sed -i "/^\[connectivity\]/a quic-port=${listenPort}" /etc/dcv/dcv.conf

          # remove AWSCLI version 1
          apt-get remove awscli -y
          apt-get autoremove -y

          cd /root/
          # https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html
          curl -s https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip -o awscliv2.zip
          unzip -q -o awscliv2.zip
          ./aws/install -b /usr/bin
          echo "export AWS_CLI_AUTO_PROMPT=on-partial" >> /home/kali/.zshrc

          # DCV update script
          cat << EoF > /home/kali/update-dcv
          #!/bin/zsh
          cd /tmp
          rm -f /tmp/nice-dcv-ubuntu2204-x86_64.tgz
          wget https://d1uj6qtbmh3dt5.cloudfront.net/nice-dcv-ubuntu2204-x86_64.tgz
          tar -xvzf nice-dcv-ubuntu2204-x86_64.tgz && cd nice-dcv-*-ubuntu2204-x86_64
          sudo dcv close-session kali
          sudo systemctl stop dcvserver dcv-virtual-session
          sudo apt-get install -y ./nice-dcv-server_*_amd64.ubuntu2204.deb 
          sudo apt-get install -y ./nice-dcv-web-viewer_*_amd64.ubuntu2204.deb 
          sudo apt-get install -y ./nice-xdcv_*_amd64.ubuntu2204.deb 
          sudo sed -i "s/^#enable-quic-frontend=true/enable-quic-frontend=true/g" /etc/dcv/dcv.conf
          sudo systemctl restart dcvserver dcv-virtual-session
          EoF
          chmod +x /home/kali/update-dcv
          chown kali:kali /home/kali/update-dcv           

          # AWS CLI update script
          cat << EoF > /home/kali/update-awscli
          #!/bin/zsh
          cd /tmp
          rm -f awscliv2.zip
          curl https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip -o awscliv2.zip
          unzip -q -o awscliv2.zip
          sudo ./aws/install --update -b /usr/bin
          EoF
          chmod +x /home/kali/update-awscli
          chown kali:kali /home/kali/update-awscli             

          # GPU drivers download script
          cat << EoF > /home/kali/download-NVIDIA-GRID-driver
          #!/bin/zsh
          clear
          echo  
          echo    NOTICE: These downloads are for GPU instances and are available to AWS customers only
          echo  
          echo By downloading, you agree to conditions and are bound by license terms as stated on 
          echo    https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/install-nvidia-driver.html
          echo 
          sudo mkdir -p /home/kali/Downloads/Drivers
          sudo chown kali:users /home/kali/Downloads/Drivers
          cd /home/kali/Downloads/Drivers
          pwd
          aws s3 cp --recursive s3://ec2-linux-nvidia-drivers/latest/ .
          chmod +x /home/kali/Downloads/Drivers/*.run
          EoF

          cat << EoF > /home/kali/download-NVIDIA-Gaming-driver
          #!/bin/zsh
          clear
          echo  
          echo    NOTICE: These downloads are for GPU instances and are available to AWS customers only
          echo  
          echo By downloading, you agree to conditions and are bound by license terms as stated on 
          echo    https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/install-nvidia-driver.html
          echo 
          sudo mkdir -p /home/kali/Downloads/Drivers
          sudo chown kali:users /home/kali/Downloads/Drivers
          cd /home/kali/Downloads/Drivers
          pwd
          aws s3 cp --recursive s3://nvidia-gaming/linux/latest/ .
          sudo mkdir -p /etc/nvidia
          sudo curl -s -o /etc/nvidia/GridSwCert.txt "https://nvidia-gaming.s3.amazonaws.com/GridSwCert-Archive/GridSwCertLinux_2021_10_2.cert"
          EoF

          cat << EoF > /home/kali/download-AMD-driver
          #!/bin/zsh
          clear
          echo  
          echo    NOTICE: These downloads are for GPU instances and are available to AWS customers only
          echo  
          echo By downloading, you agree to conditions and are bound by license terms as stated on 
          echo    https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/install-amd-driver.html
          echo 
          sudo mkdir -p /home/kali/Downloads/Drivers
          sudo chown kali:users /home/kali/Downloads/Drivers
          cd /home/kali/Downloads/Drivers
          pwd
          aws s3 cp --recursive s3://ec2-amd-linux-drivers/latest/ .
          cat << EOF > /home/kali/Downloads/Drivers/amd-xorg.conf
          Section "ServerLayout"
              Identifier     "Layout0"
              Screen          0 "Screen0"
              InputDevice     "Keyboard0" "CoreKeyboard"
              InputDevice     "Mouse0" "CorePointer"
          EndSection
          Section "Files"
              ModulePath "/opt/amdgpu/lib64/xorg/modules/drivers"
              ModulePath "/opt/amdgpu/lib/xorg/modules"
              ModulePath "/opt/amdgpu-pro/lib/xorg/modules/extensions"
              ModulePath "/opt/amdgpu-pro/lib64/xorg/modules/extensions"
              ModulePath "/usr/lib64/xorg/modules"
              ModulePath "/usr/lib/xorg/modules"
          EndSection
          Section "InputDevice"
              # generated from default
              Identifier     "Mouse0"
              Driver         "mouse"
              Option         "Protocol" "auto"
              Option         "Device" "/dev/psaux"
              Option         "Emulate3Buttons" "no"
              Option         "ZAxisMapping" "4 5"
          EndSection
          Section "InputDevice"
              # generated from default
              Identifier     "Keyboard0"
              Driver         "kbd"
          EndSection
          Section "Monitor"
              Identifier     "Monitor0"
              VendorName     "Unknown"
              ModelName      "Unknown"
          EndSection
          Section "Device"
              Identifier     "Device0"
              Driver         "amdgpu"
              VendorName     "AMD"
              BoardName      "Radeon MxGPU V520"
              BusID          "PCI:0:30:0"
          EndSection
          Section "Extensions"
              Option         "DPMS" "Disable"
          EndSection
          Section "Screen"
              Identifier     "Screen0"
              Device         "Device0"
              Monitor        "Monitor0"
              DefaultDepth   24
              Option         "AllowEmptyInitialConfiguration" "True"
              SubSection "Display"
                  Virtual    3840 2160
                  Depth      32
              EndSubSection
          EndSection
          EOF
          EoF

          chmod +x /home/kali/download-*-driver
          chown kali:kali /home/kali/download-*-driver

          # Fix "Authentication Required to Create Managed Color Device" prompt
          cat << EoF > /etc/polkit-1/localauthority/50-local.d/45-allow-colord.pkla
          [Allow Colord all Users]
          Identity=unix-user:*
          Action=org.freedesktop.color-manager.create-device;org.freedesktop.color-manager.create-profile;org.freedesktop.color-manager.delete-device;org.freedesktop.color-manager.delete-profile;org.freedesktop.color-manager.modify-device;org.freedesktop.color-manager.modify-profile
          ResultAny=no
          ResultInactive=no
          ResultActive=yes
          EoF

          # text console: DCV virtual sessions only
          systemctl isolate multi-user.target
          systemctl set-default multi-user.target

          systemctl daemon-reload
          systemctl enable --now dcvserver dcv-virtual-session.service

          # cfn-init completed so signal success or not
          python3 /usr/local/bin/cfn-signal -e $? --stack ${AWS::StackName} --resource ec2Instance --region ${AWS::Region}
      Tags:
        - Key: Name
          Value: !Ref ec2Name
        - Key: StackName
          Value: !Sub ${AWS::StackName}
        - Key: StackId
          Value: !Sub ${AWS::StackId}
        - Key: GitHub
          Value: https://github.com/aws-samples/amazon-ec2-nice-dcv-samples

  elasticIP:
    Condition: useElasticIP
    Type: AWS::EC2::EIP
    Properties:
      Domain: vpc
      InstanceId: !Ref ec2Instance
      Tags:
        - Key: StackName
          Value: !Sub ${AWS::StackName}
        - Key: StackId
          Value: !Sub ${AWS::StackId}
        - Key: Name
          Value: !Sub "[${AWS::StackName}] - ${ec2Name}"
        - Key: GitHub
          Value: https://github.com/aws-samples/amazon-ec2-nice-dcv-samples

Outputs:
  EC2Instance:
    Description: EC2 instance console
    Value: !Sub "https://${AWS::Region}.console.aws.amazon.com/ec2/home?region=${AWS::Region}#Instances:search=${ec2Instance}"

  SSMsessionManager:
    Description: SSM Session Manager login ("sudo passwd kali" to change password)
    Value: !Sub "https://${AWS::Region}.console.aws.amazon.com/systems-manager/session-manager/${ec2Instance}"

  DCVwebConsole:
    Description: DCV web console (login as kali)
    Value: !Sub
      - "https://${IpAddress}:${listenPort}"
      - IpAddress:
          !If [
            displayPublicIP,
            !GetAtt ec2Instance.PublicIp,
            !GetAtt ec2Instance.PrivateIp,
          ]