AWSTemplateFormatVersion: 2010-09-09
Parameters:
  APIKEY:
    Description: Api key for use API in Cloud One Workload Security
    Type: String
Resources:
  C1WSGroup:
    Type: AWS::IAM::Group
  C1SNSPolicy:
    Type: AWS::IAM::Policy
    Properties:
      PolicyName: C1WSUserPolicy
      PolicyDocument:
        Version: "2012-10-17"
        Statement:
        - Effect: Allow
          Action:
          - 'sns:Publish'
          Resource: '*'
      Groups: [!Ref C1WSGroup]
  C1WSUser:
    Type: AWS::IAM::User
  C1WSGroupAddUser:
    Type: AWS::IAM::UserToGroupAddition
    Properties:
      GroupName: !Ref 'C1WSGroup'
      Users: [!Ref 'C1WSUser']
  ACCESSKey:
    Type: AWS::IAM::AccessKey
    Properties: 
      UserName: !Ref C1WSUser
  LambdaRole:
    Type: 'AWS::IAM::Role'
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - lambda.amazonaws.com
            Action:
              - 'sts:AssumeRole'
      Path: /
      Policies:
        - PolicyName: c1-ws-lambdasecurityhub-role
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action: 
                  - securityhub:BatchImportFindings
                  - securityhub:BatchUpdateFindings
                  - logs:CreateLogStream
                  - logs:CreateLogGroup
                  - logs:PutLogEvents
                Resource: '*'
  C1WSLambdaSecurityHub:
    Type: AWS::Lambda::Function
    Properties:
      Handler: index.lambda_handler
      Role: 
        Fn::GetAtt:
          - LambdaRole
          - Arn
      Runtime: "python3.6"
      MemorySize: 128
      Timeout: 50
      Code:
        ZipFile: |
          import datetime
          import json
          import os
          import boto3
          from botocore.retries.standard import RetryPolicy

          def generate_finding_title(title):
              return "Trend Micro: {}".format(title)

          def verify_required_properties(ws_event):
              result = True
              required_properties = [
                  'HostOwnerID', 
                  'HostInstanceID', 
                  'TenantID', 
                  'EventID', 
                  "EventType", 
                  'LogDate', 
                  'HostAssetValue', 
                  'HostGroupID', 
                  'HostGroupName', 
                  'HostID', 
                  'Hostname', 
                  'HostSecurityPolicyID', 
                  'HostSecurityPolicyName'
                  ]
              for prop in required_properties:
                  if not prop in ws_event:
                      result = False
                      print(result)
                      return result

          def select_asff_eventType(EventType, types):
              if(EventType == 'PacketLog'):
                  types.append("Unusual Behaviors/Network Flow")
              elif(EventType == 'IntegrityEvent'):
                  types.append("Unusual Behaviors/VM")
              elif(EventType == 'LogInspectionEvent' or EventType == 'PayloadLog'):
                  types.append("Unusual Behaviors/VM")
                  types.append("Unusual Behaviors/Application")
              elif(EventType == 'WebReputationEvent' or EventType == 'AntiMalwareEvent'):
                  types.append("TPPs/Execution")
              return types

          def antimalwareStatusAction(action):
              state = ""
              if action == "Cleaned" or action == "Deleted":
                  state = "REMOVED"
              elif action == "Quarantined" or action == "Access Denied":
                  state = "OBSERVED"
              else:
                  state = "REMOVAL_FAILED"
              return state
                  
          def addAdditionalInformation(Event, finding):
              if 'SystemEvent' in Event["EventType"]:
                  pass
              if "PacketLog" in Event["EventType"]:
                  finding['Severity']['Product'] = 0
                  finding['Severity']['Normalized'] = int(20)
                  select_asff_eventType(Event["EventType"], finding["Types"])
                  finding['Title'] = "Trend Micro: Repeated attempted network connection on instance {}".format(Event['HostInstanceID'])
              if "PayloadLog" in Event["EventType"]:
                  if 'Severity' in Event:
                      finding['Severity']['Product'] = int(Event['Severity'])
                      finding['Severity']['Normalized'] = int(int(Event['Severity']) * 17.5)
                  select_asff_eventType(Event["EventType"], finding["Types"])
                  finding['Title'] = "Trend Micro: Rule [{}] triggered".format(Event['Reason'])
              if "AntiMalwareEvent" in Event["EventType"]:
                  finding['Malware'] = [
                      {
                          "Name": Event['MalwareName'],
                          "Path": Event['InfectedFilePath'],
                          "State": antimalwareStatusAction(Event['ScanResultString']),
                          }
                      ]
                  if Event['ScanResultString'] == "Cleaned" or Event['ScanResultString'] == "Deleted":
                      finding['Severity']['Label'] = "INFORMATIONAL"
                  elif Event['ScanResultString'] == "Quarantined":
                      finding['Severity']['Label'] = "LOW"
                  else: 
                      finding['Severity']['Label'] = "MEDIUM"
                  select_asff_eventType(Event["EventType"], finding["Types"])
                  finding['Title'] = "Malware [{}] detected".format(Event['MalwareName'])
              if "WebReputationEvent" in Event["EventType"]:
                  if 'Risk' in Event:
                      finding['Severity']['Product'] = int(Event['Risk'])
                      finding['Severity']['Normalized'] = int(int(Event['Risk']) * 17.5)
                  select_asff_eventType(Event["EventType"], finding["Types"])
                  finding['Title'] = "High risk web request to IP [{}]".format(Event['TargetIP'])
              if "IntegrityEvent" in Event["EventType"]:
                  if 'Severity' in Event:
                      finding['Severity']['Product'] = int(Event['Severity'])
                      finding['Severity']['Normalized'] = int(int(Event['Severity']) * 17.5)
                  select_asff_eventType(Event["EventType"], finding["Types"])
                  finding['Title'] = "Unexpected change to object [{}]".format(Event['Key'])
              if "LogInspectionEvent" in Event["EventType"]:
                  if 'OSSEC_Level' in Event:
                      finding['Severity']['Product'] = int(Event['OSSEC_Level'])
                      if int(Event['OSSEC_Level']) >= 13:
                          finding['Severity']['Normalized'] = int(int(Event['OSSEC_Level']) * 6.5)
                      else:
                          finding['Severity']['Normalized'] = int(int(Event['OSSEC_Level']) * 5)
                  select_asff_eventType(Event["EventType"], finding["Types"])
                  finding['Title'] = Event['OSSEC_Description']
              if "AppControlEvent" in Event["EventType"]:
                  Process = {
                      "Name": Event['FileName'],
                      "Path": Event['Path'],
                      "Pid": Event['ProcessID'],
                  }
                  finding["Process"] = Process
                  if(Event['Operation'] <=1):
                      finding['Severity']['Label'] = "LOW"
                      finding['ProductFields']['trend-micro:SHA256'] = Event['SHA256'] if 'SHA256' in Event else ''
                      finding['Title'] = "User {} performs allowed Execution of Unrecognized Software".format(Event['UserName'])
                  else:
                      finding['Severity']['Label'] = "INFORMATIONAL"
                      finding['ProductFields']['trend-micro:SHA256'] = Event['SHA256'] if 'SHA256' in Event else ''
                      finding['Title'] = "User {} tried performs Execution but Agent taking action of Block".format(Event['UserName'])
                  finding['Types'].append("Unusual Behaviors/Application")
                  finding['Title'] = "User {} performs allowed Execution of Unrecognized Software".format(Event['UserName'])
              if 'Tags' in Event:
                  finding['ProductFields']['trend-micro:Tags'] = Event['Tags']
              if 'OriginString' in Event:
                  finding['ProductFields']['trend-micro:Origin'] = Event['OriginString']
              return finding

          def workload_security_event_to_asff(ws_event, region, awsaccountid):
              event_types = {
                  'SystemEvent': 'system',
                  'PacketLog': 'firewall',
                  'PayloadLog': 'intrusionprevention',
                  'AntiMalwareEvent': 'antimalware',
                  'WebReputationEvent': 'webreputation',
                  'IntegrityEvent': 'integrity',
                  'LogInspectionEvent': 'log',
                  'AppControlEvent': 'applicationcontrol'
                  }

              finding = {
                  "SchemaVersion": "2018-10-08",
                  "Id": ws_event["UniqueID"],
                  "ProductArn": f"arn:aws:securityhub:{region}:{awsaccountid}:product/{awsaccountid}/default",
                  "GeneratorId": "trend-micro-workload-security-{}".format(event_types[ws_event["EventType"]]),
                  "AwsAccountId": awsaccountid,
                  "Types": [
                      ],
                  "CreatedAt": ws_event['LogDate'],
                  "UpdatedAt": datetime.datetime.utcnow().isoformat("T") + "Z",
                  "Severity": {
                      "Product": 0,
                      "Normalized": 0
                      },
                  "ProductFields": {
                      'ProviderName': "Trend Micro Cloud One",
                      "ProviderVersion": "20",
                      'trend-micro:TenantName': ws_event['TenantName'] if 'TenantName' in ws_event else '',
                      'trend-micro:TenantID': str(ws_event['TenantID']) if 'TenantID' in ws_event else '',
                      'trend-micro:EventID': str(ws_event['EventID']) if 'EventID' in ws_event else '',
                      'trend-micro:HostAssetValue': str(ws_event['HostAssetValue']) if 'HostAssetValue' in ws_event else '',
                      'trend-micro:HostGroupID': str(ws_event['HostGroupID']) if 'HostGroupID' in ws_event else '',
                      'trend-micro:HostGroupName': ws_event['HostGroupName'] if 'HostGroupName' in ws_event else '',
                      'trend-micro:HostID': str(ws_event['HostID']) if 'HostID' in ws_event else '',
                      'trend-micro:HostInstanceID': str(ws_event['HostInstanceID']) if 'HostInstanceID' in ws_event else '',
                      'trend-micro:Hostname': ws_event['Hostname'] if 'Hostname' in ws_event else '',
                      'trend-micro:HostSecurityPolicyID': str(ws_event['HostSecurityPolicyID']) if 'HostSecurityPolicyID' in ws_event else '',
                      'trend-micro:HostSecurityPolicyName': ws_event['HostSecurityPolicyName'] if 'HostSecurityPolicyName' in ws_event else '',
                      'trend-micro:Origin' : ws_event['OriginString'] if 'OriginString' in ws_event else '',
                      'trend-micro:EventType' : ws_event['EventType'] if 'EventType' in ws_event else ''
                      },
                  "Description": "Workload Security Event, type: {}".format(event_types[ws_event["EventType"]]),
                  "Resources": [
                      {
                          "Type": "AwsEc2Instance",
                          "Id":  ws_event['HostInstanceID'] if 'HostInstanceID' in ws_event else ''
                          }
                      ],
                  "Title": "Cloud One Workload Security push the following event type: {} for HostName: {}".format(event_types[ws_event["EventType"]], ws_event['Hostname'] if 'Hostname' in ws_event else ''),
                  }
              converted_event = addAdditionalInformation(ws_event, finding)
              return converted_event

          def lambda_handler(event, context):
              total_events = 0
              saved_events = 0
              securityhub = boto3.client("securityhub")
              region = boto3.session.Session().region_name
              awsaccountid = boto3.client("sts").get_caller_identity()["Account"]
              if 'Records' in event:
                  for e in event['Records']:
                      if 'EventSource' in e and e['EventSource'] == 'aws:sns':
                          print("Amazon SNS message received")
                          if 'Sns' in e:
                              ws_events = None
                              try:
                                  ws_events = json.loads(e['Sns']['Message'])
                              except Exception as err:
                                  print("Could not extract the Workload Security event(s) from the SNS message. Threw exception:\n{}".format(err))
                              aff_events = []
                              if ws_events:
                                  print("Found {} Workload Security events...processing".format(len(ws_events)))
                                  for ws_event in ws_events:
                                      """if('HostInstanceID' in ws_event):"""
                                      total_events += 1
                                      if not ws_event["EventType"] == 'SystemEvent' or verify_required_properties(ws_event):
                                          aff_event = workload_security_event_to_asff(ws_event=ws_event, region=region, awsaccountid=awsaccountid)
                                          aff_events.append(aff_event)
                                          print(aff_event)
                                      else: print("Specified event does not have the required properties to properly process it")
                              if len(aff_events) > 0:
                                  response = securityhub.batch_import_findings(Findings=aff_events)
                                  print(json.dumps(response))
                  return {
                      'total_events': total_events,
                      'saved_events': saved_events,
                      'issues': (total_events - saved_events)
                      }
  ApplicationBlockRule:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: ApplicationBlockRule
      Handler: index.lambda_handler
      Role: 
        Fn::GetAtt:
          - LambdaRole
          - Arn
      Runtime: "python3.8"
      MemorySize: 128
      Timeout: 50
      Code:
        ZipFile: |
          import json
          import os
          import urllib3
          import boto3

          http = urllib3.PoolManager()

          mainUrl = "https://cloudone.trendmicro.com/api/"
          headers = {
              'api-version': 'v1',
              'Content-Type': 'application/json',
              'api-secret-key': os.environ["APIKEY"]
          }

          def lambda_handler(event, context):
              for finding in event['detail']['findings']:
                  if(finding['GeneratorId'] == 'trend-micro-workload-security-applicationcontrol'):
                      productARN = finding['ProductArn']
                      findingId = finding['Id']
                      hostID = finding['ProductFields']['trend-micro:HostID']
                      processHash = finding['ProductFields']['trend-micro:SHA256']
                      searchSoftwareChanges(hostID, processHash, findingId, productARN)

          def searchSoftwareChanges(computerID, sha256, findingId, productARN):
              url=mainUrl+"softwarechanges/search"
              print("HostID", computerID)
              print("SHA256", sha256)
              payload = json.dumps({"searchCriteria": [{"fieldName": "computerID","numericTest": "equal","numericValue": computerID},{"fieldName": "sha256","stringTest": "equal","stringValue": sha256}]})

              r = http.request("POST", url, headers=headers, body=payload)
              response = json.loads(r.data.decode('utf-8'))
              if(len(response['softwareChanges'])>0):
                  changeProcessID = response['softwareChanges'][0]['ID']
                  blockExecution(changeProcessID)
                  updateFinding(findingId, productARN)
              else:
                  if(getAgentConfiguraton(computerID, sha256)):
                      updateFinding(findingId, productARN)

          def blockExecution(changeProcessID):
              url = mainUrl+"softwarechanges/review"
              
              payload = json.dumps({
                  "softwareChangeIDs": [
                      changeProcessID,
                  ],
                  "action": "block"
              })
              r = http.request("POST", url, headers=headers, body=payload)
              response = json.loads(r.data.decode('utf-8'))
              print(response)
              return True

          def updateFinding(findingId, productARN):
              securityhub = boto3.client("securityhub")
              response = securityhub.batch_update_findings(
                  FindingIdentifiers=[{
                      'Id': findingId,
                      'ProductArn': productARN
                  }], 
                  Severity={ 
                      'Label': 'INFORMATIONAL', 
                  }
              ) 
              print(response)

          def getAgentConfiguraton(computerID, sha256):
              url = mainUrl+"computers/search?expand=applicationControl"
              payload = json.dumps({
                  "searchCriteria": [{"fieldName": "computerID","idValue": computerID,"idTest": "equal"}]})
              r = http.request("POST", url, headers=headers, body=payload)
              response = json.loads(r.data.decode('utf-8'))
              rulesetID = response["computers"][0]['applicationControl']['rulesetID']
              urlR = mainUrl+"rulesets/{}/rules/search".format(rulesetID)
              payloadR = json.dumps({
                  "searchCriteria": [{"fieldName": "sha256","stringTest": "equal","stringValue": sha256},{"fieldName": "action","choiceTest": "not-equal","choiceValue": "block"}]
              })
              r = http.request("POST", urlR, headers=headers, body=payloadR)
              response = json.loads(r.data.decode('utf-8'))
              if(len(response['applicationControlRules']) > 0):
                  ruleID = response['ID']
                  url=mainUrl+"rulesets/{}/rules/{}".format(rulesetID, ruleID)
                  payload=json.dumps({"action": "block"})
                  r = http.request("POST", url, headers=headers, body=payload)
                  response = json.loads(r.data.decode('utf-8'))
                  print(response)
                  return True
              else:
                  return False
      Environment:
        Variables:
          APIKEY: !Ref APIKEY
  SNSC1WS:
    Type: AWS::SNS::Topic
    Properties:
      DisplayName: C1WorkloadSecuritySNS
      TopicName: C1WorkloadSecuritySNS
  Subscription:
    Type: AWS::SNS::Subscription
    Properties:
      Endpoint: !GetAtt
        - C1WSLambdaSecurityHub
        - Arn
      Protocol: lambda
      TopicArn: !Ref SNSC1WS
    DependsOn: [ "C1WSLambdaSecurityHub" ]
  InvokeLambdaPermission:
    Type: AWS::Lambda::Permission
    Properties:
      Action: lambda:InvokeFunction
      Principal: sns.amazonaws.com
      SourceArn: !Ref SNSC1WS
      FunctionName: !GetAtt
        - C1WSLambdaSecurityHub
        - Arn
Outputs:
  SNSTopic:
    Description: The SNS ARN to Workload Security
    Value: !Ref SNSC1WS
  AccessKey:
    Description: The Access key for Worload Security User
    Value: !Ref ACCESSKey
  SecretKey:
    Description: The Secret key for Worload Security User
    Value: !GetAtt ACCESSKey.SecretAccessKey
  LambdaFunction:
    Description: Lamba Function to send SNS Notification to Security Hub
    Value: !Ref C1WSLambdaSecurityHub
  CustomActionFunction:
    Description: Lamba Function to perform custom Action with Cloud One
    Value: !Ref ApplicationBlockRule