AWSTemplateFormatVersion: '2010-09-09'
Description: >-
  Connects this AWS account to Velobyte. Velobyte can see which GPU instance
  types this account can launch, file quota increase requests, and launch, stop
  and terminate model endpoint instances that it tags velobyte=1, and admit
  HTTPS from Cloudflare to its own security group. It cannot read or change
  anything else: no data, no other instances, no IAM. Delete this stack to
  disconnect.

Parameters:
  ExternalId:
    Type: String
    Description: Connection ID from your Velobyte console. Do not change it.
    AllowedPattern: '[0-9a-f]{32}'
  TrustedAccountId:
    Type: String
    Description: The Velobyte AWS account allowed to assume this role.
    AllowedPattern: '[0-9]{12}'
  CallbackUrl:
    Type: String
    Default: https://velobyte.ai/models/api/aws/callback
    Description: Where the stack tells Velobyte it is ready.

Resources:
  VelobyteAccessRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: !Sub 'VelobyteAccess-${AWS::StackName}'
      Description: Velobyte model endpoints (tag-scoped EC2 only)
      MaxSessionDuration: 3600
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              AWS: !Sub 'arn:aws:iam::${TrustedAccountId}:root'
            Action: sts:AssumeRole
            Condition:
              StringEquals:
                sts:ExternalId: !Ref ExternalId
      Policies:
        - PolicyName: velobyte-endpoints
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Sid: SeeGpuCapacity
                Effect: Allow
                Action:
                  - ec2:DescribeInstanceTypeOfferings
                  - ec2:DescribeInstanceTypes
                  - ec2:DescribeRegions
                  - ec2:DescribeInstances
                  - ec2:DescribeImages
                  - ec2:DescribeSecurityGroups
                  - ec2:DescribeSubnets
                  - ec2:DescribeVpcs
                  - ec2:DescribeAvailabilityZones
                  - servicequotas:GetServiceQuota
                  - servicequotas:GetAWSDefaultServiceQuota
                  - servicequotas:ListServiceQuotas
                  - servicequotas:ListRequestedServiceQuotaChangeHistory
                  - servicequotas:ListRequestedServiceQuotaChangeHistoryByQuota
                Resource: '*'
              - Sid: LaunchOnlyTaggedInstancesAndVolumes
                Effect: Allow
                Action: ec2:RunInstances
                Resource:
                  - !Sub 'arn:${AWS::Partition}:ec2:*:${AWS::AccountId}:instance/*'
                  - !Sub 'arn:${AWS::Partition}:ec2:*:${AWS::AccountId}:volume/*'
                Condition:
                  StringEquals:
                    aws:RequestTag/velobyte: '1'
              - Sid: LaunchSupportingResources
                Effect: Allow
                Action: ec2:RunInstances
                Resource:
                  - !Sub 'arn:${AWS::Partition}:ec2:*::image/*'
                  - !Sub 'arn:${AWS::Partition}:ec2:*:${AWS::AccountId}:subnet/*'
                  - !Sub 'arn:${AWS::Partition}:ec2:*:${AWS::AccountId}:network-interface/*'
                  - !Sub 'arn:${AWS::Partition}:ec2:*:${AWS::AccountId}:security-group/*'
              - Sid: TagOnlyAtCreation
                Effect: Allow
                Action: ec2:CreateTags
                Resource:
                  - !Sub 'arn:${AWS::Partition}:ec2:*:${AWS::AccountId}:instance/*'
                  - !Sub 'arn:${AWS::Partition}:ec2:*:${AWS::AccountId}:volume/*'
                  - !Sub 'arn:${AWS::Partition}:ec2:*:${AWS::AccountId}:security-group/*'
                Condition:
                  StringEquals:
                    ec2:CreateAction:
                      - RunInstances
                      - CreateSecurityGroup
              - Sid: ManageOnlyTaggedInstances
                Effect: Allow
                Action:
                  - ec2:StopInstances
                  - ec2:StartInstances
                  - ec2:TerminateInstances
                Resource: !Sub 'arn:${AWS::Partition}:ec2:*:${AWS::AccountId}:instance/*'
                Condition:
                  StringEquals:
                    aws:ResourceTag/velobyte: '1'
              - Sid: CreateTaggedSecurityGroupWithNoInboundRules
                Effect: Allow
                Action: ec2:CreateSecurityGroup
                Resource: !Sub 'arn:${AWS::Partition}:ec2:*:${AWS::AccountId}:security-group/*'
                Condition:
                  StringEquals:
                    aws:RequestTag/velobyte: '1'
              - Sid: CreateSecurityGroupInVpc
                Effect: Allow
                Action: ec2:CreateSecurityGroup
                Resource: !Sub 'arn:${AWS::Partition}:ec2:*:${AWS::AccountId}:vpc/*'
              - Sid: AdmitHttpsOnTaggedGroupOnly
                Effect: Allow
                Action: ec2:AuthorizeSecurityGroupIngress
                Resource: !Sub 'arn:${AWS::Partition}:ec2:*:${AWS::AccountId}:security-group/*'
                Condition:
                  StringEquals:
                    aws:ResourceTag/velobyte: '1'
              - Sid: ReadPublicAmiParameters
                Effect: Allow
                Action:
                  - ssm:GetParameter
                  - ssm:GetParameters
                Resource: !Sub 'arn:${AWS::Partition}:ssm:*::parameter/aws/service/*'
              - Sid: FileGpuQuotaRequests
                Effect: Allow
                Action: servicequotas:RequestServiceQuotaIncrease
                Resource: '*'
              - Sid: QuotaRequestsNeedServiceLinkedRole
                Effect: Allow
                Action: iam:CreateServiceLinkedRole
                Resource: !Sub 'arn:${AWS::Partition}:iam::${AWS::AccountId}:role/aws-service-role/servicequotas.amazonaws.com/*'
                Condition:
                  StringEquals:
                    iam:AWSServiceName: servicequotas.amazonaws.com

  # Tells Velobyte the role exists so the console connects without copy and paste.
  NotifyRole:
    Type: AWS::IAM::Role
    Properties:
      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'

  NotifyFunction:
    Type: AWS::Lambda::Function
    Properties:
      Runtime: python3.12
      Handler: index.handler
      Timeout: 30
      Role: !GetAtt NotifyRole.Arn
      Code:
        ZipFile: |
          import json, urllib.request
          def send(url, body, method='POST'):
              data = json.dumps(body).encode()
              headers = {'Content-Type': '' if method == 'PUT' else 'application/json', 'User-Agent': 'VelobyteConnect/1.0 (+https://velobyte.ai/models)'}
              req = urllib.request.Request(url, data=data, method=method, headers=headers)
              return urllib.request.urlopen(req, timeout=10)
          def handler(event, context):
              props = event['ResourceProperties']
              try:
                  send(props['CallbackUrl'], {'event': event['RequestType'], 'externalId': props['ExternalId'],
                                              'roleArn': props['RoleArn'], 'awsAccountId': props['AccountId']})
              except Exception as error:
                  print('callback failed', error)
              # Never block stack create or delete on the callback.
              send(event['ResponseURL'], {'Status': 'SUCCESS', 'PhysicalResourceId': props['RoleArn'],
                   'StackId': event['StackId'], 'RequestId': event['RequestId'], 'LogicalResourceId': event['LogicalResourceId'],
                   'Reason': 'ok'}, method='PUT')

  Notify:
    Type: Custom::VelobyteNotify
    Properties:
      ServiceToken: !GetAtt NotifyFunction.Arn
      Version: '2'
      CallbackUrl: !Ref CallbackUrl
      ExternalId: !Ref ExternalId
      RoleArn: !GetAtt VelobyteAccessRole.Arn
      AccountId: !Ref AWS::AccountId

Outputs:
  RoleArn:
    Description: The role Velobyte uses. Delete this stack to disconnect.
    Value: !GetAtt VelobyteAccessRole.Arn