Armorblox sdk code changes (demisto#20901) (demisto#20902)

* Playbook, Classifier, Docker tag changes * armorblox-sdk code changes with devdemisto docker image * Updated code changes for armorblox-sdk version 0.1.2 * Updated code changes for get_page_token * updated code review changes in Armorblox.py * updated code with demisto debug command for integration * Updated Docker image to demisto/armorblox:1.0.0.33173 * resolved linting errors * resolved linting errors * Updated the code as per comments * Added Release Notes file(1_0_4.md) with new docker image * - Improved the RN - Bumped the pack's version * little code improvements * cheating coverage * remove comment for cheating coverage * cheating coverage * cheating coverage - changed the location of the pragma: no coverage comment Co-authored-by: Ankita Sharma <[email protected]> Co-authored-by: SamhithaTatipalli <[email protected]> Co-authored-by: SamhithaTatipalli <[email protected]> Co-authored-by: Ankita Sharma <[email protected]> Co-authored-by: Shachar Kidor <[email protected]> Co-authored-by: Ankita Sharma <[email protected]> Co-authored-by: SamhithaTatipalli <[email protected]> Co-authored-by: SamhithaTatipalli <[email protected]> Co-authored-by: Ankita Sharma <[email protected]>
humio · Sep 1, 2022 · b62c9b9 · b62c9b9
1 parent f3910d5
commit b62c9b9
Show file tree

Hide file tree

Showing 8 changed files with 224 additions and 304 deletions.
diff --git a/Packs/Armorblox/.pack-ignore b/Packs/Armorblox/.pack-ignore
@@ -2,4 +2,9 @@
 ignore=IM111
 
 [file:classifier-Armorblox_-_Classifier.json]
-ignore=BA101
+ignore=BA101
+
+[known_words]
+Armorblox
+armorblox-sdk
+Classifiers
diff --git a/Packs/Armorblox/Classifiers/classifier-Armorblox_-_Classifier.json b/Packs/Armorblox/Classifiers/classifier-Armorblox_-_Classifier.json
@@ -4,57 +4,32 @@
     "feed": false,
     "id": "e8ded555-9409-4d33-842c-45d29b6ab31c",
     "keyTypeMap": {
-        "Abuse Report": "Armorblox Abuse Mailbox Report",
-        "Extortion": "Armorblox Inbound Threat",
-        "Graymail": "Armorblox Inbound Threat",
-        "Impersonation: Employee": "Armorblox Inbound Threat",
-        "Impersonation: VIP": "Armorblox Inbound Threat",
-        "Impersonation: VIP (Requesting Gift Card)": "Armorblox Inbound Threat",
-        "Impersonation:Employee": "Armorblox Inbound Threat",
-        "Impersonation:VIP": "Armorblox Inbound Threat",
-        "PCI Bank Account Number": "Armorblox Outbound Threat",
-        "PCI Credit Card Number": "Armorblox Outbound Threat",
-        "PCI IBAN": "Armorblox Outbound Threat",
-        "PCI Routing Number": "Armorblox Outbound Threat",
-        "PII Passport": "Armorblox Outbound Threat",
-        "PII Social Security Number": "Armorblox Outbound Threat",
-        "PII Tax Number": "Armorblox Outbound Threat",
-        "Passwords": "Armorblox Outbound Threat",
-        "Payment Fraud (External)": "Armorblox Inbound Threat",
-        "Payment Fraud (Internal)": "Armorblox Inbound Threat",
-        "Payroll Fraud": "Armorblox Inbound Threat",
-        "Phish URL (Attachment)": "Armorblox Inbound Threat",
-        "Phish URL (Mail Body)": "Armorblox Inbound Threat",
-        "Potential Account Compromise": "Armorblox Inbound Threat",
-        "Ransomware": "Armorblox Inbound Threat",
-        "Social Engineering": "Armorblox Inbound Threat"
+        "ABUSE_INCIDENT_TYPE": "Armorblox Abuse Mailbox Report",
+        "DLP_INCIDENT_TYPE": "Armorblox Outbound Threat",
+        "THREAT_INCIDENT_TYPE": "Armorblox Inbound Threat"
     },
     "name": "Armorblox - Classifier",
     "transformer": {
         "complex": {
-            "accessor": "",
+            "accessor": "incident_type",
             "filters": [],
-            "root": "policy_names",
+            "root": "incidents",
             "transformers": [
                 {
                     "args": {
                         "descending": {
-                            "isContext": false,
                             "value": {
-                                "complex": null,
                                 "simple": "false"
                             }
                         }
                     },
                     "operator": "sort"
                 },
                 {
-                    "args": {},
                     "operator": "FirstArrayElement"
                 }
             ]
-        },
-        "simple": ""
+        }
     },
     "type": "classification",
     "version": -1,

diff --git a/Packs/Armorblox/Integrations/Armorblox/Armorblox.py b/Packs/Armorblox/Integrations/Armorblox/Armorblox.py
@@ -1,4 +1,7 @@
+from typing import Any
+
 import demistomock as demisto  # noqa: F401
+from armorblox.client import Client as AbxBaseClient
 from CommonServerPython import *  # noqa: F401
 import dateparser
 import requests
@@ -17,53 +20,38 @@
 API_KEY = demisto.params().get('apikey')
 verify_certificate = not demisto.params().get('insecure', False)
 proxy = demisto.params().get('proxy', False)
-BASE_URL = f"https://{TENANT_NAME}.armorblox.io/api/v1beta1/organizations/{TENANT_NAME}"
-
-payload: Dict = {}
-headers = {
-    'x-ab-authorization': f'{API_KEY}'
-}
 
 
-class Client(BaseClient):
+class Client(AbxBaseClient):
     """Client class to interact with the service API
     This Client implements API calls, and does not contain any Demisto logic.
     Should only do requests and return data.
-    It inherits from BaseClient defined in CommonServer Python.
-    Most calls use _http_request() that handles proxy, SSL verification, etc.
     """
 
-    def get_incidents(self, orderBy="ASC", pageSize=None, pageToken=None, first_fetch=None) -> List[Dict[str, Any]]:
-        request_params: Dict[str, Any] = {}
+    def get_incidents(self, orderBy='ASC', pageSize=None, pageToken=None, first_fetch=None):
+        request_params = {'orderBy': orderBy}
 
-        request_params['orderBy'] = orderBy
         if pageToken == -1 and first_fetch:
             request_params['timeFilter'] = first_fetch
         elif pageToken and first_fetch:
             request_params['timeFilter'] = first_fetch
             request_params['pageToken'] = pageToken
+
         if pageSize:
             request_params['pageSize'] = pageSize
-        return self._http_request(
-            method='GET',
-            url_suffix='/incidents',
-            params=request_params
-        )
+
+        response_json, next_page_token, total_count = self.incidents.list(params=request_params)
+        return response_json, next_page_token
 
     def get_incident_details(self, incident_id):
-        request_params: Dict[str, Any] = {}
-        return self._http_request(
-            method='GET',
-            url_suffix='/incidents/{}'.format(incident_id),
-            params=request_params
-        )
+        return self.incidents.get(incident_id)
 
 
 def makehash():
     return collections.defaultdict(makehash)
 
 
-def test_module(client: Client) -> str:
+def test_module(client: Client) -> str:  # pragma: no coverage
     """Tests API connectivity and authentication'
     Returning 'ok' indicates that the integration works like it is supposed to.
     Connection to the service is successful.
@@ -86,27 +74,16 @@ def test_module(client: Client) -> str:
     return 'ok'
 
 
-def get_page_token(client, pageToken=None):
-    response = client.get_incidents(pageSize=MAX_INCIDENTS_TO_FETCH, pageToken=pageToken, first_fetch=FIRST_FETCH)
-    if 'next_page_token' in response.keys():
-        return response['next_page_token']
-    else:
-        return None
-
-
 def get_incidents_list(client, pageToken, first_fetch):
     """
     Hits the Armorblox API and returns the list of fetched incidents.
     """
-    response = client.get_incidents(pageSize=MAX_INCIDENTS_TO_FETCH, pageToken=pageToken, first_fetch=first_fetch)
-    results = []
-    if 'incidents' in response.keys():
-        results = response['incidents']
-
+    results, next_page_token = client.get_incidents(pageSize=MAX_INCIDENTS_TO_FETCH, pageToken=pageToken,
+                                                    first_fetch=first_fetch)
     # For each incident, get the details and extract the message_id
     for result in results:
         result['message_ids'] = get_incident_message_ids(client, result["id"])
-    return results
+    return results, next_page_token
 
 
 def get_incident_message_ids(client, incident_id):
@@ -152,17 +129,17 @@ def fetch_incidents_command(client):
     start_time: Any
     # pageToken fetched from demisto lastRun
     pageToken = int()
-    response = {}
     incidents = []
     if 'start_time' not in last_run.keys():
         pageToken = -1
-        response = client.get_incidents(pageSize=1, pageToken=pageToken, first_fetch=FIRST_FETCH)
-        if 'incidents' in response.keys():
-            start_time = response['incidents'][0]['date']
+        response, next_page_token = client.get_incidents(pageSize=1, pageToken=pageToken, first_fetch=FIRST_FETCH)
+        if response:
+            response = response[0]
+            start_time = response.get('date')
             start_time = dateparser.parse(start_time)
-            message_ids = get_incident_message_ids(client, response['incidents'][0]['id'])
-            response['incidents'][0]['message_ids'] = message_ids
-            curr_incident = {'rawJSON': json.dumps(response['incidents'][0]), 'details': json.dumps(response['incidents'][0])}
+            message_ids = get_incident_message_ids(client, response.get('id'))
+            response['message_ids'] = message_ids
+            curr_incident = {'rawJSON': json.dumps(response), 'details': json.dumps(response)}
             incidents.append(curr_incident)
 
     if last_run and 'pageToken' in last_run.keys():
@@ -172,18 +149,16 @@ def fetch_incidents_command(client):
         start_time = dateparser.parse(last_run.get('start_time'))
 
     start_time = start_time.timestamp()
-    incidents_data = get_incidents_list(client, pageToken=pageToken, first_fetch=FIRST_FETCH)
-    pageToken = get_page_token(client, pageToken=pageToken)
+    incidents_data, pageToken = get_incidents_list(client, pageToken=pageToken, first_fetch=FIRST_FETCH)
     last_time = start_time
 
     for incident in incidents_data:
-        dt = incident['date']
+        dt = incident.get('date')
         parsed_date = dateparser.parse(dt)
         assert parsed_date is not None, f'failed parsing {dt}'
-        dt = parsed_date.timestamp()
+        dt = int(parsed_date.timestamp())
         # Update last run and add incident if the incident is newer than last fetch
-        if dt > start_time:
-
+        if dt > int(start_time):
             curr_incident = {'rawJSON': json.dumps(incident), 'details': json.dumps(incident)}
             last_time = dt
             incidents.append(curr_incident)
@@ -192,17 +167,15 @@ def fetch_incidents_command(client):
     return incidents
 
 
-def main():
+def main():  # pragma: no coverage
     ''' EXECUTION '''
-    LOG('command is %s' % (demisto.command(), ))
+    demisto.info(f'Command being called is {demisto.command()}')
     try:
 
         client = Client(
-            base_url=BASE_URL,
-            verify=verify_certificate,
-            headers=headers,
-            proxy=proxy)
-
+            api_key=API_KEY,
+            instance_name=TENANT_NAME
+        )
         if demisto.command() == "fetch-incidents":
             incident_results = fetch_incidents_command(client)
             demisto.incidents(incident_results)

diff --git a/Packs/Armorblox/Integrations/Armorblox/Armorblox.yml b/Packs/Armorblox/Integrations/Armorblox/Armorblox.yml
@@ -73,7 +73,7 @@ script:
     - contextPath: Armorblox.Threat.remediation_actions
       description: Should be the remediation action name for the incident under inspection
       type: string
-  dockerimage: demisto/python3:3.10.5.31928
+  dockerimage: demisto/armorblox:1.0.0.33173
   isfetch: true
   runonce: false
   script: ''