Azure AD Admin - 17046 - Refactored code in risk detection trigger (#…

…2586)

plugins/azure_ad_admin/.CHECKSUM

@@ -1,5 +1,5 @@
 {
-	"spec": "127260067e5296b22e8c5dd975448768",
+	"spec": "65941e2ec4ae777adfa6f97a970ac8c5",
 	"manifest": "00a0692089d7ca3cbf498f443004b42d",
 	"setup": "16febaa8fc47f2dde2ca771e5e88d209",
 	"schemas": [

plugins/azure_ad_admin/Dockerfile

@@ -1,4 +1,4 @@
-FROM --platform=linux/amd64 rapid7/insightconnect-python-3-slim-plugin:5.5.2
+FROM --platform=linux/amd64 rapid7/insightconnect-python-3-slim-plugin:5.5.3
 
 LABEL organization=rapid7
 LABEL sdk=python

plugins/azure_ad_admin/help.md

@@ -1056,7 +1056,7 @@ Trigger `risk_detection` needs Application permission to set as `IdentityRiskEve
 
 # Version History
 
-* 4.1.2 - Updated SDK to the latest version | Added additional details in requirements section
+* 4.1.2 - Updated SDK to the latest version | Added additional details in requirements section | `Risk Detection`: Fixed issue where detections were triggered randomly
 * 4.1.1 - Update requirements in help.md
 * 4.1.0 - New actions Enable Device, Disable Device, Get Device, Search Device, Delete Device
 * 4.0.0 - Get User Info action: fix data validation | New action: Change User Password

plugins/azure_ad_admin/icon_azure_ad_admin/triggers/risk_detection/trigger.py

@@ -1,12 +1,21 @@
+import json
+import time
+from datetime import datetime
+from typing import Any, Dict, List
+
 import insightconnect_plugin_runtime
-from .schema import RiskDetectionInput, RiskDetectionOutput, Input, Output, Component
+import requests
+from dateutil import parser
+from insightconnect_plugin_runtime.exceptions import PluginException
+from insightconnect_plugin_runtime.helper import clean
 
 # Custom imports below
 from icon_azure_ad_admin.util.komand_clean_with_nulls import remove_null_and_clean
-from insightconnect_plugin_runtime.exceptions import PluginException
-import time
-import requests
-import json
+
+from .schema import Component, Input, Output, RiskDetectionInput, RiskDetectionOutput
+
+DEFAULT_REQUESTS_TIMEOUT = 30
+DETECTED_RISK_DATE_FIELD = "detectedDateTime"
 
 
 class RiskDetection(insightconnect_plugin_runtime.Trigger):
@@ -17,67 +26,123 @@ def __init__(self):
             input=RiskDetectionInput(),
             output=RiskDetectionOutput(),
         )
-        self.risk_level = ""
-        self.found = {}
 
-    def initialize(self):
-        new_risks = self.get_risks()
+    def run(self, params={}):
+        # START INPUT BINDING - DO NOT REMOVE - ANY INPUTS BELOW WILL UPDATE WITH YOUR PLUGIN SPEC AFTER REGENERATION
+        risk_level = params.get(Input.RISK_LEVEL)
+        frequency = params.get(Input.FREQUENCY, 60)
+        # END INPUT BINDING - DO NOT REMOVE
 
-        try:
-            result = new_risks["value"]
-        except KeyError:
-            raise PluginException(
-                cause="Unexpected output format.",
-                assistance="The output from Azure Active Directory was not in the expected format. Please contact support for help.",
-                data=new_risks,
+        # Initialize the trigger with starting data point
+        current_risks = self._get_risks(risk_level)
+        latest_risk_detection_time = self._get_latest_risk_detection_time(current_risks)
+
+        while True:
+            # Retrieve current 'riskDetection' data
+            self.logger.info("Retrieving risks data from the API...")
+            current_risks = self._get_risks(risk_level, latest_risk_detection_time)
+
+            # Filter all results by their 'detectedDateTime', returning only results whose time is greater than
+            # the last risk detection 'detectedDateTime'.
+            new_risks = list(
+                filter(
+                    lambda element: self.parse_datetime(element.get(DETECTED_RISK_DATE_FIELD))
+                    > latest_risk_detection_time,
+                    current_risks,
+                )
             )
 
-        for risk in result:
-            self.found[risk.get("id")] = True
+            if new_risks:
+                self.logger.info(f"Found new {len(new_risks)} risks. Returning...")
+                latest_risk_detection_time = self._get_latest_risk_detection_time(new_risks)
+                for risk in new_risks:
+                    self.send({Output.RISK: risk})
+            else:
+                self.logger.info(f"No new risks found. Sleeping for {frequency} seconds...")
+            time.sleep(frequency)
 
-    def get_risks(self):
-        headers = self.connection.get_headers(self.connection.get_auth_token())
+    @staticmethod
+    def parse_datetime(input_datetime_str: str) -> datetime:
+        """
+        Parse a string representation of a datetime object and return a :class:`datetime.datetime` object.
+        This function accepts a string formatted according to a standard datetime format, primarily supporting ISO 8601.
 
-        if self.risk_level and self.risk_level != "all":
-            risk_detect_endpoint = f"https://graph.microsoft.com/beta/{self.connection.tenant}/riskDetections?$filter=riskLevel eq '{self.risk_level}'"
-        else:
-            risk_detect_endpoint = f"https://graph.microsoft.com/beta/{self.connection.tenant}/riskDetections"
+        :param input_datetime_str: The string representation of the datetime to be parsed.
+        :type: str
 
-        new_risks = requests.get(risk_detect_endpoint, headers=headers)
+        :returns: A datetime object representing the parsed datetime.
+        :rtype: datetime
+        """
 
-        if not new_risks.status_code == 200:
+        return parser.parse(input_datetime_str, ignoretz=True)
+
+    def _get_latest_risk_detection_time(self, risk_detections: List[Dict[str, Any]]) -> datetime:
+        """
+        Retrieve the latest risk detection time from a list of risk detections.
+
+        :param risk_detections: A list of risk detection records, where each record is a dictionary containing details about the risk detection event.
+        :type: List[Dict[str, Any]]
+
+        :return: The latest risk detection time found in the provided list.
+        :rtype: datetime
+        """
+
+        if not risk_detections:
+            return datetime.utcnow()
+        return self.parse_datetime(risk_detections[-1].get(DETECTED_RISK_DATE_FIELD))
+
+    def _get_risks(self, risk_level: str, latest_risk_detection_time: datetime = None) -> List[Dict[str, Any]]:
+        """
+        Retrieve the risk detection event from MS Graph API.
+
+        :param risk_level: Risk detection level for results to be filtered on.
+        :type: str
+
+        :param latest_risk_detection_time: Determines whether API should return all results greater than specific time or not. Defaults to None.
+        :type: datetime
+
+        :return: The latest risk detection time found in the provided list.
+        :rtype: List[Dict[str, Any]]
+        """
+
+        # The 'riskLevel' filter condition to be applied when it's not 'all'
+        risk_filter = f" and riskLevel eq '{risk_level}'" if risk_level != "all" else ""
+
+        # Sending request to the MS Graph API with detectedDateTime filter, it allows to return only new detections
+        # Setup $top to 500 as it's the maximum number of record that can be returned using that endpoint
+        response = requests.get(
+            "https://graph.microsoft.com/v1.0/identityProtection/riskDetections",
+            headers=self.connection.get_headers(self.connection.get_auth_token()),
+            params=clean(
+                {
+                    "$filter": f"{DETECTED_RISK_DATE_FIELD} gt {latest_risk_detection_time.isoformat()}Z" + risk_filter
+                    if latest_risk_detection_time
+                    else "",
+                    "$top": 500,
+                }
+            ),
+            timeout=DEFAULT_REQUESTS_TIMEOUT,
+        )
+
+        if not response.status_code == 200:
             raise PluginException(
-                cause=f"Risk Detections returned an unexpected response: {new_risks.status_code}",
+                cause=f"Risk Detections returned an unexpected response: {response.status_code}",
                 assistance="Please contact support for help.",
-                data=new_risks.text,
+                data=response.text,
             )
 
         try:
-            return new_risks.json()
-        except json.decoder.JSONDecodeError:
-            raise PluginException(preset=PluginException.Preset.INVALID_JSON, data=new_risks)
-
-    def poll(self):
-        new_risks = self.get_risks()
-
-        try:
-            result = remove_null_and_clean(new_risks["value"])
+            # Cleaning out the response records and sorting them out manually in ascending order
+            # because MS Graph API was not ordering them properly
+            return sorted(
+                remove_null_and_clean(response.json()["value"]),
+                key=lambda element: self.parse_datetime(element.get(DETECTED_RISK_DATE_FIELD)),
+            )
         except KeyError:
             raise PluginException(
                 cause="Unexpected output format.",
                 assistance="The output from Azure Active Directory was not in the expected format. Please contact support for help.",
-                data=new_risks,
+                data=response.text,
             )
-
-        for risk in result:
-            if risk.get("id") not in self.found:
-                self.found[risk.get("id")] = True
-                self.send({Output.RISK: risk})
-
-    def run(self, params={}):
-        self.risk_level = params.get(Input.RISK_LEVEL)
-        self.frequency = params.get(Input.FREQUENCY, 60)
-        self.initialize()
-        while True:
-            self.poll()
-            time.sleep(self.frequency)
+        except json.decoder.JSONDecodeError as error:
+            raise PluginException(preset=PluginException.Preset.INVALID_JSON, data=error)

plugins/azure_ad_admin/plugin.spec.yaml

@@ -13,7 +13,7 @@ supported_versions: ["2022-05-30"]
 status: []
 sdk:
   type: slim
-  version: 5.5.2
+  version: 5.5.3
   user: nobody
 key_features:
   - Add and remove users
@@ -35,7 +35,7 @@ requirements:
   Browse and select your application and click `Add`."
 troubleshooting: "Trigger `risk_detection` needs Application permission to set as `IdentityRiskEvent.Read.All`"
 version_history:
-  - "4.1.2 - Updated SDK to the latest version | Added additional details in requirements section"
+  - "4.1.2 - Updated SDK to the latest version | Added additional details in requirements section | `Risk Detection`: Fixed issue where detections were triggered randomly"
   - "4.1.1 - Update requirements in help.md"
   - "4.1.0 - New actions Enable Device, Disable Device, Get Device, Search Device, Delete Device"
   - "4.0.0 - Get User Info action: fix data validation | New action: Change User Password"

plugins/azure_ad_admin/requirements.txt

@@ -1,4 +1,5 @@
 # List third-party dependencies here, separated by newlines.
 # All dependencies must be version-pinned, eg. requests==1.2.0
 # See: https://pip.pypa.io/en/stable/user_guide/#requirements-files
+python-dateutil==2.9.0.post0
 parameterized==0.8.1