# this cloudformation is provided as a supplement to the Blog Post on "How to run your Jamulus Server on AWS".
# It is offered under the GPL3 licence, see http://www.gnu.org/licenses/gpl.html as-is, with no warranty or claim to functionality or liability resulting from the use thereof.
# Make sure to read and understand this Cloudformation stack before you use it!
# Author: Helge Aufderheide (AWS Professional Services)

Parameters:
  # This parameters allows to seperate multple deployments (if you need them)
  ProjectName:
    Default: "JamulusStack"
    Description: "a name for your stack, used as a prefix to your resources"
    Type: String
  MaxUsers:
    Default: 16
    Description: "maximum number of users you want to allow. If >25, consider increasing the server size"
    Type: Number


Resources:
  # We need a cluster to run our tasks is. This is just an empty shell, no actual machines will be created
  FargateCluster:
    Type: "AWS::ECS::Cluster"
    Properties:
      ClusterName: !Sub ${ProjectName}-cluster

  # We create a loggroup where the cluster can store its logs and reduce retention to 1 day.
  FargateLogGroup:
    Type: "AWS::Logs::LogGroup"
    Properties: 
      LogGroupName: !Sub ${ProjectName}-loggroup
      RetentionInDays: 1
  
  # This is the actual Jamulus Task we want to run
  TaskDefinition:
    Type: AWS::ECS::TaskDefinition
    Properties:
      Family: !Sub "${ProjectName}"
      # These two paramters determine the size of our Server. Increase it if you experience crashes or have many participants (see https://aws.amazon.com/fargate/pricing/ for pricing)
      Cpu: 1024
      Memory: 2048
      NetworkMode: awsvpc
      RequiresCompatibilities:
        - FARGATE
      ExecutionRoleArn: !Ref TaskExecutionRole
      TaskRoleArn: !Ref TaskRole
      ContainerDefinitions:
        - Name: jamulus
          Image: grundic/jamulus
          PortMappings:
            - ContainerPort: 22124
              Protocol: 'udp'
#              HostPort: 22124 is where Jamulus listenes to users
          EntryPoint:
            - "Jamulus"
            - "--server"
            - "--nogui"
            - "--numchannels"
            - !Sub "${MaxUsers}"
          Environment:
            - Name: TZ
              Value: America/Los_Angeles
          LogConfiguration:
            LogDriver: awslogs
            Options:
              awslogs-group: !Ref FargateLogGroup
              awslogs-region: !Ref 'AWS::Region'
              awslogs-stream-prefix: 'access'

  # We need a VPC and subnets. If you already haave them setup remove this portion
  VPC:
    Type: AWS::EC2::VPC
    Properties:
      EnableDnsSupport: true
      EnableDnsHostnames: true
      CidrBlock: '10.0.0.0/16'

  PublicSubnetOne:
    Type: AWS::EC2::Subnet
    Properties:
      AvailabilityZone: !Select 
        - 0
        - !GetAZs 
          Ref: 'AWS::Region'
      VpcId: !Ref VPC
      CidrBlock: '10.0.0.0/24'
      MapPublicIpOnLaunch: true
  
  PublicSubnetTwo:
    Type: AWS::EC2::Subnet
    Properties:
      AvailabilityZone: !Select 
        - 1
        - !GetAZs 
          Ref: 'AWS::Region'
      VpcId: !Ref 'VPC'
      CidrBlock: '10.0.1.0/24'
      MapPublicIpOnLaunch: true

  # Internet Routing, so that Jamulus can pull its image from dockerhub and can be reached   
  InternetGateway:
    Type: AWS::EC2::InternetGateway

  GatewayAttachement:
    Type: AWS::EC2::VPCGatewayAttachment
    Properties:
      VpcId: !Ref 'VPC'
      InternetGatewayId: !Ref 'InternetGateway'

  PublicRouteTable:
    Type: AWS::EC2::RouteTable
    Properties:
      VpcId: !Ref 'VPC'

  PublicRoute:
    Type: AWS::EC2::Route
    DependsOn: GatewayAttachement
    Properties:
      RouteTableId: !Ref 'PublicRouteTable'
      DestinationCidrBlock: '0.0.0.0/0'
      GatewayId: !Ref 'InternetGateway'

  PublicSubnetOneRouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      SubnetId: !Ref PublicSubnetOne
      RouteTableId: !Ref PublicRouteTable

  PublicSubnetTwoRouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      SubnetId: !Ref PublicSubnetTwo
      RouteTableId: !Ref PublicRouteTable

  # The security group of the Jamulus task requires ingress on 22124 (from users) and egress (to dockerhub)
  JamulusSecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties: 
      GroupName: !Sub "fargate-${ProjectName}"
      GroupDescription: 'A security group to allow ingress to ${ProjectName} hosts'
      SecurityGroupIngress:
        - CidrIp: 0.0.0.0/0
          Description: 'allows udp on 22124 from everywhere'
          FromPort: 22124
          IpProtocol: udp
          ToPort: 22124
      SecurityGroupEgress:
        - CidrIp: 0.0.0.0/0
          Description: 'allows udp on 22124 from everywhere'
          FromPort: 443
          IpProtocol: tcp
          ToPort: 443
      VpcId: !Ref VPC
  
  # The task needs a role that can pull images and do logging        
  TaskExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Statement:
        - Effect: Allow
          Principal:
            Service: [ecs-tasks.amazonaws.com]
          Action: ['sts:AssumeRole']
      Path: /
      Policies:
        - PolicyName: AmazonECSTaskExecutionRolePolicy
          PolicyDocument:
            Statement:
            - Effect: Allow
              Action:
                - 'ecr:GetAuthorizationToken'
                - 'ecr:BatchCheckLayerAvailability'
                - 'ecr:GetDownloadUrlForLayer'
                - 'ecr:BatchGetImage'
                - 'logs:CreateLogStream'
                - 'logs:PutLogEvents'
                - 'cloudwatch:PutMetricData'
                - 'cloudwatch:PutMetricStream'
              Resource: '*'

  TaskRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Statement:
        - Effect: Allow
          Principal:
            Service: [ecs-tasks.amazonaws.com]
          Action: ['sts:AssumeRole']
      Path: /

  # This is the ECS Manager Lambda function. The role can start the Jamulus task, ask for their IP and send Email (for advanced usage)
  LambdaExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service:
              - lambda.amazonaws.com
            Action:
              - sts:AssumeRole
      Path: "/"
      Policies:
        - PolicyName: "LambdaManagerInlineRights"
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - logs:CreateLogGroup
                  - logs:CreateLogStream
                  - logs:PutLogEvents
                Resource: arn:aws:logs:*:*:*
              - Effect: Allow
                Action:
                  - ecs:Describe*
                  - ecs:List*
                  - ecs:RunTask
                  - ecs:StartTask
                  - ecs:StopTask
                Condition:
                  ArnEquals:
                      ecs:cluster: !GetAtt FargateCluster.Arn
                Resource: "*"
              - Effect: Allow
                Action:
                  - ec2:DescribeNetworkInterfaces
                Resource: "*"
              - Effect: Allow
                Action:
                  - iam:PassRole
                Resource: !Sub arn:aws:iam::${AWS::AccountId}:role/${AWS::StackName}-*

  # The lambda function itself
  # Important Environment Variables: 
  # - VpcId: The ID of the VPC to run the task in
  # - SubnetIds: (comma-seperated) list of subnets
  # - Security Groups: (comma-seperated) list of security group ids
  # - Project Name: The Name of the project, so we can identify Tasks that belong to it
  ManagerLambda:
    Type: AWS::Lambda::Function
    Properties:  
      Environment:
        Variables:
          VpcId: !Ref VPC
          SubnetIds: !Sub "${PublicSubnetOne},${PublicSubnetOne}"
          SecurityGroups: !Sub "${JamulusSecurityGroup}"
          ProjectName: !Sub "${ProjectName}"
      FunctionName: !Sub "${ProjectName}-ECSManager"
      Handler: index.lambda_handler
      Role: 
        Fn::GetAtt:
          - LambdaExecutionRole
          - Arn
      Runtime: python3.8
      Timeout: 10
      Code:
        ZipFile: |
          import json
          import boto3
          import os

          # Important Environment Variables: 
          # - VpcId: The ID of the VPC to run the task in
          # - SubnetIds: (comma-seperated) list of subnets
          # - Security Groups: (comma-seperated) list of security group ids
          # - Project Name: The Name of the project, so we can identify Tasks that belong to it

          ProjectName = os.environ['ProjectName']
          VpcId = os.environ['VpcId']
          SubnetIds = os.environ['SubnetIds'].split(",")
          SecurityGroups = os.environ.get('SecurityGroups').split(",")
          Task = ProjectName

          # Clients needed: ECS for Task access, EC2 for network configuration (IP address)
          ECS = boto3.client('ecs')
          EC2 = boto3.client('ec2')

          def lambda_handler(event, context):
              # We expect event["ACTION"] = START|STOP|CHECK to tell the lambda what to do
              if not 'ACTION' in event:
                  return {'statusCode': 400, 'body': json.dumps('No ACTION provided')}
                  
              if not event['ACTION'] in ["START","STOP","CHECK"]:
                  return {'statusCode': 400, 'body': json.dumps('No valid ACTION provided: START, STOP, CHECK')}
              
              ## We are good to see what is running
              current_tasks = ECS.list_tasks(cluster = ProjectName+"-cluster", family = ProjectName).get('taskArns',[])

              if len(current_tasks)>1:
                  return {'statusCode': 400,  'body': json.dumps('More than 1 task running... please stop them manually') }
              else:
                  print("Found the following tasks:",current_tasks)
              
              # We have none of some tasks running: lets do what we are supposed to
              
              if event['ACTION'] ==  "START":
                  if len(current_tasks)>0:
                      return { 'statusCode': 400,'body': json.dumps('Task already running, if it is stopping, please retry later')}
                  print("Starting task")    
                  response = ECS.run_task(
                      cluster=ProjectName+"-cluster",
                      count=1,
                      launchType='FARGATE',
                      networkConfiguration={
                          'awsvpcConfiguration': {
                              'subnets': SubnetIds,
                              'securityGroups': SecurityGroups,
                              'assignPublicIp': 'ENABLED'
                          }},
                      taskDefinition=Task
                  )
                  return { 'statusCode': 200,'body': json.dumps(response,default=str)}

              elif event['ACTION'] == "STOP":
                  response = "NOTHING DONE"
                  for task_tostop in current_tasks:
                      print("Stopping",task_tostop)
                      response = ECS.stop_task(
                          cluster=ProjectName+"-cluster",
                          task=task_tostop
                      )
                  return { 'statusCode': 200,'body': json.dumps(response,default=str)}
                  
              
              else:
                  print("DescribingTasks")
                  if len(current_tasks) ==0:
                      return { 'statusCode': 200,'body': json.dumps("No Tasks running")}
                  response = ECS.describe_tasks(
                      cluster=ProjectName+"-cluster",
                      tasks=current_tasks
                          )

                  eni_ids = [det['value'] for det in  response['tasks'][0]['attachments'][0]['details'] if det['name']=="networkInterfaceId"]
                  nw_config = EC2.describe_network_interfaces( NetworkInterfaceIds=eni_ids )    
                  return { 'statusCode': 200,'body': json.dumps({"STATUS": response['tasks'][0]['containers'][0]['lastStatus'],'public_IP':nw_config["NetworkInterfaces"][0]["Association"]} ,default=str)}
        
    
    
