Armis AlertsActivities and Device Data Connector Enhancements

Solutions/Armis/Data Connectors/ArmisAlertsActivities/ArmisAlertActivitySentinelConnector/consts.py

@@ -25,11 +25,13 @@
 CONNECTION_STRING = os.environ.get("AzureWebJobsStorage", "")
 ARMIS_ALERTS_TABLE = os.environ.get("ArmisAlertsTableName", "")
 ARMIS_ACTIVITIES_TABLE = os.environ.get("ArmisActivitiesTableName", "")
-IS_AVOID_DUPLICATES = os.environ.get("AvoidDuplicates", "")
 WORKSPACE_ID = os.environ.get("WorkspaceID", "")
 WORKSPACE_KEY = os.environ.get("WorkspaceKey", "")
 CHUNK_SIZE = 35
 FILE_SHARE = "funcstatemarkershare"
-CHECKPOINT_FILE = "funcarmisalertsfile"
+CHECKPOINT_FILE_TIME = "funcarmisalertsfile"
+CHECKPOINT_FILE_OFFSET = "armisalertoffset"
 LOG_FORMAT = "Armis Alerts Activities Connector: (method = {}) : {}"
 REQUEST_TIMEOUT = 300
+CHECKPOINT_TABLE_NAME = "ArmisAlertActivityCheckpoint"
+FUNCTION_APP_TIMEOUT_SECONDS = 570

Solutions/Armis/Data Connectors/ArmisAlertsActivities/ArmisAlertActivitySentinelConnector/exports_store.py

@@ -0,0 +1,91 @@
+import logging
+
+from azure.data.tables import TableClient, UpdateMode
+from azure.core.exceptions import ResourceNotFoundError, ResourceExistsError, HttpResponseError
+
+
+class ExportsTableStore:
+
+    def __init__(self, connection_string, table_name):
+        self.connection_string = connection_string
+        self.table_name = table_name
+
+    def create(self):
+        with TableClient.from_connection_string(self.connection_string, self.table_name) as table_client:
+            try:
+                table_client.create_table()
+                logging.info("Checkpoint Table created")
+            except ResourceExistsError:
+                logging.warning("Checkpoint Table already exists")
+
+    def post(self, pk: str, rk: str, data: dict = None):
+        with TableClient.from_connection_string(self.connection_string, self.table_name) as table_client:
+            entity_template = {
+                "PartitionKey": pk,
+                "RowKey": rk,
+            }
+            if data is not None:
+                entity_template.update(data)
+            try:
+                table_client.create_entity(entity_template)
+            except Exception as e:
+                logging.warning("could not post entity to table")
+                logging.warning(e)
+                raise e
+
+    def get(self, pk: str, rk: str):
+        with TableClient.from_connection_string(self.connection_string, self.table_name) as table_client:
+            try:
+                logging.info(
+                    "looking for {} - {} on table {}".format(pk, rk, self.table_name))
+                return table_client.get_entity(pk, rk)
+            except ResourceNotFoundError:
+                return None
+
+    def upsert(self, pk: str, rk: str, data: dict = None):
+        with TableClient.from_connection_string(self.connection_string, self.table_name) as table_client:
+            logging.info("upserting {} - {} on table {}".format(pk, rk, self.table_name))
+            entity_template = {
+                "PartitionKey": pk,
+                "RowKey": rk,
+            }
+            if data is not None:
+                entity_template.update(data)
+            return table_client.upsert_entity(mode=UpdateMode.REPLACE, entity=entity_template)
+
+    def update_if_found(self, pk: str, rk: str, data: dict = None):
+        if self.get(pk, rk) is not None:
+            self.merge(pk, rk, data)
+
+    def query_by_partition_key(self, pk):
+        table_client = TableClient.from_connection_string(
+            self.connection_string, self.table_name)
+        parameters = {u"key": pk}
+        name_filter = u"PartitionKey eq @key"
+        try:
+            return table_client.query_entities(name_filter, parameters=parameters)
+        except HttpResponseError as e:
+            print(e.message)
+            return []
+
+    def batch(self, operations):
+        with TableClient.from_connection_string(self.connection_string, self.table_name) as table_client:
+            return table_client.submit_transaction(operations=operations)
+
+    def list_all(self):
+        table_client = TableClient.from_connection_string(
+            self.connection_string, self.table_name)
+        return table_client.list_entities()
+
+    def merge(self, pk: str, rk: str, data: dict = None):
+        with TableClient.from_connection_string(self.connection_string, self.table_name) as table_client:
+            logging.info("upserting {} - {} on table {}".format(pk, rk, self.table_name))
+            entity_template = {
+                "PartitionKey": pk,
+                "RowKey": rk,
+            }
+            if data is not None:
+                entity_template.update(data)
+            return table_client.upsert_entity(mode=UpdateMode.MERGE, entity=entity_template)
+
+

Solutions/Armis/Data Connectors/ArmisAlertsActivities/ArmisAlertActivitySentinelConnector/state_manager.py

@@ -36,3 +36,13 @@ def get(self):
             return self.file_cli.download_file().readall().decode()
         except ResourceNotFoundError:
             return None
+
+    def delete(self):
+        """Delete method for deleting the data from Azure Storage.
+
+        This method will delete the file from Azure Storage.
+        """
+        try:
+            self.file_cli.delete_file()
+        except ResourceNotFoundError:
+            raise ResourceNotFoundError("File not found to be deleted.")

Solutions/Armis/Data Connectors/ArmisAlertsActivities/ArmisAlertActivitySentinelConnector/utils.py

@@ -22,15 +22,14 @@ def __init__(self) -> None:
                 {"WorkspaceKey": consts.WORKSPACE_KEY},
                 {"ArmisSecretKey": consts.API_KEY},
                 {"AzureWebJobsStorage": consts.CONNECTION_STRING},
-                {"AvoidDuplicates": consts.IS_AVOID_DUPLICATES},
                 {"ArmisAlertsTableName": consts.ARMIS_ALERTS_TABLE},
                 {"ArmisActivitiesTableName": consts.ARMIS_ACTIVITIES_TABLE},
             ]
         )
         self._secret_key = consts.API_KEY
         self.get_access_token()
         self.state_manager_obj = StateManager(
-            connection_string=consts.CONNECTION_STRING, file_path=consts.CHECKPOINT_FILE
+            connection_string=consts.CONNECTION_STRING, file_path=consts.CHECKPOINT_FILE_TIME
         )
 
     def check_environment_var_exist(self, environment_var):
@@ -78,10 +77,10 @@ def make_rest_call(self, method, url, params=None, headers=None, data=None, retr
         """
         __method_name = inspect.currentframe().f_code.co_name
         try:
-            response = requests.request(
-                method, url, headers=self.header, params=params, data=data, timeout=consts.REQUEST_TIMEOUT
-            )
             for _ in range(retry_401 + 1):
+                response = requests.request(
+                    method, url, headers=self.header, params=params, data=data, timeout=consts.REQUEST_TIMEOUT
+                )
                 if response.status_code == 200:
                     response_json = response.json()
                     logging.info(
@@ -222,40 +221,14 @@ def get_formatted_time(self, alert_time):
             )
             raise ArmisException()
 
-    def post_alert_checkpoint(self, alert):
-        """Post alert checkpoint.
-
-        Args:
-            alert (dict): last alert from data
-        """
-        __method_name = inspect.currentframe().f_code.co_name
-        try:
-            alert_time = self.get_formatted_time(alert["time"][:19])
-            self.state_manager_obj.post(str(alert_time))
-            logging.info(
-                consts.LOG_FORMAT.format(__method_name, "Alerts checkpoint updated : {}.".format(str(alert_time)))
-            )
-        except KeyError as err:
-            logging.error(consts.LOG_FORMAT.format(__method_name, "Key error : {}.".format(err)))
-            raise ArmisException()
-
-        except ArmisException:
-            raise ArmisException()
-
-        except Exception as err:
-            logging.error(
-                consts.LOG_FORMAT.format(__method_name, "Error while posting alerts checkpoint : {}.".format(err))
-            )
-            raise ArmisException()
-
     def get_access_token(self):
         """get_access_token method will fetch the access token using api and set it in header for further use."""
         __method_name = inspect.currentframe().f_code.co_name
         try:
             body = {"secret_key": self._secret_key}
             logging.info(consts.LOG_FORMAT.format(__method_name, "Getting access token."))
             response = self.make_rest_call(method="POST", url=consts.URL + consts.ACCESS_TOKEN_SUFFIX, data=body)
-            access_token = response["data"]["access_token"]
+            access_token = response.get("data", {}).get("access_token")
             self.header.update({"Authorization": access_token})
             logging.info(consts.LOG_FORMAT.format(__method_name, "Generated access token Successfully."))
         except KeyError as err:

Solutions/Armis/Data Connectors/ArmisAlertsActivities/Exceptions/ArmisExceptions.py

@@ -11,3 +11,9 @@ class ArmisDataNotFoundException(Exception):
     """ArmisDataNotFoundException class will inherit Exception class."""
 
     pass
+
+
+class ArmisTimeOutException(Exception):
+    """ArmisTimeOutException class will inherit Exception class."""
+
+    pass

Solutions/Armis/Data Connectors/ArmisAlertsActivities/azuredeploy_Connector_ArmisAlertsActivitiesAPI_AzureFunction.json

@@ -55,13 +55,6 @@
                 "description": "Enter a valid Quartz Cron-Expression (Example: 0 0 0 * * *)"
             }
         },
-        "AvoidDuplicates":{
-            "type": "bool",
-            "defaultValue": true,
-            "metadata": {
-                "description": "Indicates whether the system should avoid processing duplicate entries. Default is 'true'"
-            }
-        },
         "AppInsightsWorkspaceResourceID": {
             "type": "string",
             "metadata": {
@@ -203,7 +196,6 @@
                         "ArmisAlertsTableName": "[parameters('ArmisAlertsTableName')]",
                         "ArmisActivitiesTableName": "[parameters('ArmisActivitiesTableName')]",
                         "Schedule": "[parameters('ArmisSchedule')]",
-                        "AvoidDuplicates": "[parameters('AvoidDuplicates')]",
                         "WEBSITE_RUN_FROM_PACKAGE": "https://aka.ms/sentinel-ArmisAlertsActivities-functionapp"
                     }
                 }

Solutions/Armis/Data Connectors/ArmisAlertsActivities/requirements.txt

@@ -4,4 +4,5 @@
 
 azure-functions
 azure-storage-file-share==12.3.0
-requests
+requests
+azure-data-tables==12.1.0