Fix pre-commit linting failures in TI provider

msticpy/context/tiproviders/azure_sent_byoti.py

@@ -55,7 +55,10 @@ class AzSTI(KqlTIProvider):
     _QUERIES["linux_path"] = _QUERIES["windows_path"]
     _QUERIES["hostname"] = _QUERIES["dns"]
 
-    _REQUIRED_TABLES: ClassVar[list[str]] = ["ThreatIntelligenceIndicator"]
+    _REQUIRED_TABLES: ClassVar[list[str]] = ["ThreatIntelIndicators"]
+
+    # Confidence threshold for high severity classification
+    _HIGH_CONFIDENCE_THRESHOLD: ClassVar[int] = 80
 
     def parse_results(self: Self, response: dict) -> tuple[bool, ResultSeverity, Any]:
         """
@@ -80,19 +83,120 @@ def parse_results(self: Self, response: dict) -> tuple[bool, ResultSeverity, Any
         severity = ResultSeverity.warning
         # if this is a series (single row) return a dictionary
         if isinstance(response["RawResult"], pd.Series):
-            extracted_data: dict[str, Any] = response["RawResult"][
-                ["Action", "ThreatType", "ThreatSeverity", "Active", "ConfidenceScore"]
-            ].to_dict()
+            result_series = response["RawResult"]
+            extracted_data: dict[str, Any] = {}
+
+            # Try to use old schema fields first (for backward compatibility)
+            if "ThreatType" in result_series.index:
+                extracted_data["ThreatType"] = result_series["ThreatType"]
+            elif "Data" in result_series.index and pd.notna(result_series["Data"]):
+                # Extract from new schema Data column (STIX format)
+                stix_data = result_series["Data"]
+                if isinstance(stix_data, dict):
+                    labels = stix_data.get("labels", [])
+                    extracted_data["ThreatType"] = labels[0] if labels else "unknown"
+                else:
+                    extracted_data["ThreatType"] = "unknown"
+            else:
+                extracted_data["ThreatType"] = "unknown"
+
+            # Handle Description field
+            if "Description" in result_series.index:
+                extracted_data["Description"] = result_series["Description"]
+            elif "Data" in result_series.index and pd.notna(result_series["Data"]):
+                stix_data = result_series["Data"]
+                if isinstance(stix_data, dict):
+                    extracted_data["Description"] = stix_data.get("description", "")
+                else:
+                    extracted_data["Description"] = ""
+            else:
+                extracted_data["Description"] = ""
+
+            # Handle ConfidenceScore field (old schema) or Confidence (new schema)
+            if "ConfidenceScore" in result_series.index:
+                extracted_data["ConfidenceScore"] = result_series["ConfidenceScore"]
+            elif "Confidence" in result_series.index:
+                extracted_data["ConfidenceScore"] = result_series["Confidence"]
+            else:
+                extracted_data["ConfidenceScore"] = 0
+
+            # Handle Active field (old schema) or IsActive (new schema)
+            if "Active" in result_series.index:
+                extracted_data["Active"] = result_series["Active"]
+            elif "IsActive" in result_series.index:
+                extracted_data["Active"] = result_series["IsActive"]
+            else:
+                extracted_data["Active"] = False
+
+            # Handle Action field (old schema only, default for new schema)
+            if "Action" in result_series.index:
+                extracted_data["Action"] = result_series["Action"]
+            else:
+                extracted_data["Action"] = "alert"
+
+            # Handle ThreatSeverity field (old schema only)
+            if "ThreatSeverity" in result_series.index:
+                extracted_data["ThreatSeverity"] = result_series["ThreatSeverity"]
+            else:
+                extracted_data["ThreatSeverity"] = "unknown"
+
+            # Determine severity
             if extracted_data["Action"].lower() in ["alert", "block"]:
                 severity = ResultSeverity.high
-            return True, ResultSeverity.warning, extracted_data
+            elif extracted_data["ConfidenceScore"] >= self._HIGH_CONFIDENCE_THRESHOLD:
+                severity = ResultSeverity.high
+
+            return True, severity, extracted_data
+
         # if this is a dataframe (multiple rows)
         # concatenate the values for each column/record into a list
         # and return as a dictionary
         if isinstance(response["RawResult"], pd.DataFrame):
-            d_frame: pd.DataFrame = response["RawResult"]
+            d_frame: pd.DataFrame = response["RawResult"].copy()
+
+            # Handle ThreatType field
+            if "ThreatType" not in d_frame.columns:
+                if "Data" in d_frame.columns:
+                    d_frame["ThreatType"] = d_frame["Data"].apply(
+                        lambda x: x.get("labels", ["unknown"])[0] if isinstance(x, dict) and x.get("labels") else "unknown"
+                    )
+                else:
+                    d_frame["ThreatType"] = "unknown"
+
+            # Handle Description field
+            if "Description" not in d_frame.columns:
+                if "Data" in d_frame.columns:
+                    d_frame["Description"] = d_frame["Data"].apply(
+                        lambda x: x.get("description", "") if isinstance(x, dict) else ""
+                    )
+                else:
+                    d_frame["Description"] = ""
+
+            # Handle ConfidenceScore field
+            if "ConfidenceScore" not in d_frame.columns and "Confidence" in d_frame.columns:
+                d_frame["ConfidenceScore"] = d_frame["Confidence"]
+            elif "ConfidenceScore" not in d_frame.columns:
+                d_frame["ConfidenceScore"] = 0
+
+            # Handle Active field
+            if "Active" not in d_frame.columns and "IsActive" in d_frame.columns:
+                d_frame["Active"] = d_frame["IsActive"]
+            elif "Active" not in d_frame.columns:
+                d_frame["Active"] = False
+
+            # Handle Action field
+            if "Action" not in d_frame.columns:
+                d_frame["Action"] = "alert"
+
+            # Handle ThreatSeverity field
+            if "ThreatSeverity" not in d_frame.columns:
+                d_frame["ThreatSeverity"] = "unknown"
+
+            # Determine severity
             if d_frame["Action"].str.lower().isin(["alert", "block"]).any():
                 severity = ResultSeverity.high
+            elif "ConfidenceScore" in d_frame.columns and (d_frame["ConfidenceScore"] >= self._HIGH_CONFIDENCE_THRESHOLD).any():
+                severity = ResultSeverity.high
 
             return (
                 True,
@@ -112,6 +216,43 @@ def parse_results(self: Self, response: dict) -> tuple[bool, ResultSeverity, Any
     def _get_detail_summary(data_result: pd.DataFrame) -> pd.Series:
         # For the input frame return details in a series with
         # Details in dict
+        # Handle both old and new schema fields
+        data_result = data_result.copy()
+
+        # Handle ThreatType field
+        if "ThreatType" not in data_result.columns and "Data" in data_result.columns:
+            data_result["ThreatType"] = data_result["Data"].apply(
+                lambda x: x.get("labels", ["unknown"])[0] if isinstance(x, dict) and x.get("labels") else "unknown"
+            )
+        elif "ThreatType" not in data_result.columns:
+            data_result["ThreatType"] = "unknown"
+
+        # Handle Description field
+        if "Description" not in data_result.columns and "Data" in data_result.columns:
+            data_result["Description"] = data_result["Data"].apply(
+                lambda x: x.get("description", "") if isinstance(x, dict) else ""
+            )
+        elif "Description" not in data_result.columns:
+            data_result["Description"] = ""
+
+        # Handle ConfidenceScore field
+        if "ConfidenceScore" not in data_result.columns and "Confidence" in data_result.columns:
+            data_result["ConfidenceScore"] = data_result["Confidence"]
+        elif "ConfidenceScore" not in data_result.columns:
+            data_result["ConfidenceScore"] = 0
+
+        # Handle Active field
+        if "Active" not in data_result.columns and "IsActive" in data_result.columns:
+            data_result["Active"] = data_result["IsActive"]
+        elif "Active" not in data_result.columns:
+            data_result["Active"] = False
+
+        # Set defaults for missing fields
+        if "Action" not in data_result.columns:
+            data_result["Action"] = "alert"
+        if "ThreatSeverity" not in data_result.columns:
+            data_result["ThreatSeverity"] = "unknown"
+
         return data_result.apply(
             lambda x: {
                 "Action": x.Action,
@@ -127,10 +268,26 @@ def _get_detail_summary(data_result: pd.DataFrame) -> pd.Series:
     @staticmethod
     def _get_severity(data_result: pd.DataFrame) -> pd.Series:
         # For the input frame return severity in a series
+        data_result = data_result.copy()
+
+        # Map Confidence to ConfidenceScore if needed
+        if "Confidence" in data_result.columns and "ConfidenceScore" not in data_result.columns:
+            data_result["ConfidenceScore"] = data_result["Confidence"]
+        elif "ConfidenceScore" not in data_result.columns:
+            data_result["ConfidenceScore"] = 0
+
+        # Set default Action if not present
+        if "Action" not in data_result.columns:
+            data_result["Action"] = "alert"
+
+        # Use class constant for threshold
+        high_threshold = AzSTI._HIGH_CONFIDENCE_THRESHOLD
+
         return data_result.apply(
             lambda x: (
                 ResultSeverity.high.name
-                if x.Action.lower() in ["alert", "block"]
+                if (hasattr(x, "Action") and x.Action.lower() in ["alert", "block"]) 
+                   or (hasattr(x, "ConfidenceScore") and x.ConfidenceScore >= high_threshold)
                 else ResultSeverity.warning.name
             ),
             axis=1,

msticpy/data/queries/mssentinel/kql_sent_threatintel.yaml

@@ -11,7 +11,7 @@ defaults:
       table:
         description: Table name
         type: str
-        default: 'ThreatIntelligenceIndicator'
+        default: 'ThreatIntelIndicators'
       query_project:
         description: Project clause to limit/change return column names
         type: str
@@ -36,7 +36,7 @@ sources:
         {query_project}
         | where TimeGenerated >= datetime({start})
         | where TimeGenerated <= datetime({end})
-        | summarize arg_max(TimeGenerated, *) by IndicatorId
+        | summarize arg_max(TimeGenerated, *) by Id
         {add_query_items}'
       uri: None
   list_indicators_by_ip:
@@ -48,20 +48,10 @@ sources:
         {query_project}
         | where TimeGenerated >= datetime({start})
         | where TimeGenerated <= datetime({end})
-        | where NetworkIP in ({ip_address_list})
-          or NetworkDestinationIP in ({ip_address_list})
-          or NetworkSourceIP in ({ip_address_list})
-          or EmailSourceIpAddress in ({ip_address_list})
-        | extend IoC = tolower(
-            iff(isnotempty(NetworkIP) and NetworkIP in ({ip_address_list}), NetworkIP,
-              iff(isnotempty(NetworkDestinationIP) and NetworkDestinationIP in ({ip_address_list}), NetworkDestinationIP,
-                iff(isnotempty(NetworkSourceIP) and NetworkSourceIP in ({ip_address_list}), NetworkSourceIP,
-                  iff(isnotempty(EmailSourceIpAddress) and EmailSourceIpAddress in ({ip_address_list}), EmailSourceIpAddress, "")
-                  )
-                )
-              )
-            )
-        | summarize arg_max(TimeGenerated, *) by IndicatorId
+        | where ObservableKey in ("ipv4-addr:value", "ipv6-addr:value")
+          and ObservableValue in ({ip_address_list})
+        | extend IoC = tolower(ObservableValue)
+        | summarize arg_max(TimeGenerated, *) by Id
         {add_query_items}'
     parameters:
       ip_address_list:
@@ -78,9 +68,10 @@ sources:
         {query_project}
         | where TimeGenerated >= datetime({start})
         | where TimeGenerated <= datetime({end})
-        | where FileHashValue in~ ({file_hash_list})
-        | extend IoC = tolower(FileHashValue)
-        | summarize arg_max(TimeGenerated, *) by IndicatorId
+        | where ObservableKey in ("file:hashes.MD5", "file:hashes.SHA-1", "file:hashes.SHA-256", "file:hashes.''SHA-1''", "file:hashes.''SHA-256''")
+          and ObservableValue in~ ({file_hash_list})
+        | extend IoC = tolower(ObservableValue)
+        | summarize arg_max(TimeGenerated, *) by Id
         {add_query_items}'
     parameters:
       file_hash_list:
@@ -97,9 +88,10 @@ sources:
         {query_project}
         | where TimeGenerated >= datetime({start})
         | where TimeGenerated <= datetime({end})
-        | where FilePath in~ ({observables})
-        | extend IoC = tolower(FilePath)
-        | summarize arg_max(TimeGenerated, *) by IndicatorId
+        | where ObservableKey in ("file:name", "directory:path")
+          and ObservableValue in~ ({observables})
+        | extend IoC = tolower(ObservableValue)
+        | summarize arg_max(TimeGenerated, *) by Id
         {add_query_items}'
     parameters:
       observables:
@@ -114,14 +106,10 @@ sources:
         {query_project}
         | where TimeGenerated >= datetime({start})
         | where TimeGenerated <= datetime({end})
-        | where DomainName in~ ({domain_list})
-          or EmailSourceDomain in~ ({domain_list})
-        | extend IoC = tolower(
-            iff(isnotempty(DomainName) and DomainName in~ ({domain_list}), DomainName,
-              iff(isnotempty(EmailSourceDomain) and EmailSourceDomain in~ ({domain_list}), EmailSourceDomain, "")
-              )
-            )
-        | summarize arg_max(TimeGenerated, *) by IndicatorId
+        | where ObservableKey == "domain-name:value"
+          and ObservableValue in~ ({domain_list})
+        | extend IoC = tolower(ObservableValue)
+        | summarize arg_max(TimeGenerated, *) by Id
         {add_query_items}'
     parameters:
       domain_list:
@@ -138,14 +126,10 @@ sources:
         {query_project}
         | where TimeGenerated >= datetime({start})
         | where TimeGenerated <= datetime({end})
-        | where EmailRecipient in~ ({observables})
-          or EmailSenderAddress in~ ({observables})
-        | extend IoC = tolower(
-            iff(isnotempty(EmailRecipient) and EmailRecipient in~ ({observables}), EmailRecipient,
-              iff(isnotempty(EmailSenderAddress) and EmailSenderAddress in~ ({observables}), EmailSenderAddress, "")
-              )
-            )
-        | summarize arg_max(TimeGenerated, *) by IndicatorId
+        | where ObservableKey == "email-addr:value"
+          and ObservableValue in~ ({observables})
+        | extend IoC = tolower(ObservableValue)
+        | summarize arg_max(TimeGenerated, *) by Id
         {add_query_items}'
     parameters:
       observables:
@@ -160,9 +144,10 @@ sources:
         {query_project}
         | where TimeGenerated >= datetime({start})
         | where TimeGenerated <= datetime({end})
-        | where Url in~ ({url_list})
-        | extend IoC = tolower(Url)
-        | summarize arg_max(TimeGenerated, *) by IndicatorId
+        | where ObservableKey == "url:value"
+          and ObservableValue in~ ({url_list})
+        | extend IoC = tolower(ObservableValue)
+        | summarize arg_max(TimeGenerated, *) by Id
         {add_query_items}'
     parameters:
       url_list:

tests/context/test_tiprovider_kql.py

@@ -38,14 +38,30 @@ def __init__(self, connection_str: str = None, **kwargs):
 
         self._loaded = True
         self._connected = True
-        self._schema: Dict[str, Any] = {"ThreatIntelligenceIndicator": {}}
+        self._schema: Dict[str, Any] = {"ThreatIntelIndicators": {}}
 
         indicator_file = Path(_TEST_DATA).joinpath("as_threatintel")
         self.test_df = pd.read_pickle(indicator_file)
+
+        # Map old schema to new schema for testing
+        # Create ObservableValue and ObservableKey from old schema
         self.ip_df = self.test_df[self.test_df["NetworkIP"].str.len() > 0].copy()
+        self.ip_df["ObservableKey"] = "ipv4-addr:value"
+        self.ip_df["ObservableValue"] = self.ip_df["NetworkIP"]
         self.ip_df["IoC"] = self.ip_df["NetworkIP"].str.lower()
+
         self.url_df = self.test_df[self.test_df["Url"].str.len() > 0].copy()
+        self.url_df["ObservableKey"] = "url:value"
+        self.url_df["ObservableValue"] = self.url_df["Url"]
         self.url_df["IoC"] = self.url_df["Url"].str.lower()
+
+        # Map old fields to new schema fields
+        if "Active" in self.test_df.columns:
+            self.test_df["IsActive"] = self.test_df["Active"]
+        if "ConfidenceScore" in self.test_df.columns:
+            self.test_df["Confidence"] = self.test_df["ConfidenceScore"]
+        if "IndicatorId" in self.test_df.columns:
+            self.test_df["Id"] = self.test_df["IndicatorId"]
 
     def connect(self, connection_str: Optional[str] = None, **kwargs):
         """Mock connect function."""
@@ -64,10 +80,22 @@ def query(
         del query_source, kwargs
 
         query_toks = [tok.lower() for tok in query.split("'") if tok != ","]
+
+        # Handle new schema queries with ObservableKey and ObservableValue
+        if 'ObservableKey in ("ipv4-addr:value"' in query or 'ObservableKey in ("ipv4-addr:value", "ipv6-addr:value")' in query:
+            result_df = self.ip_df[self.ip_df["IoC"].isin(query_toks)]
+            return result_df
+
+        # Legacy support for old queries (in case they're still used)
         if "where NetworkIP" in query:
             result_df = self.ip_df[self.ip_df["IoC"].isin(query_toks)]
             return result_df
 
+        if 'ObservableKey == "url:value"' in query:
+            result_df = self.url_df[self.url_df["IoC"].isin(query_toks)]
+            return result_df
+
+        # Legacy support for old queries
         if "where Url" in query:
             result_df = self.url_df[self.url_df["IoC"].isin(query_toks)]
             return result_df
@@ -212,7 +240,7 @@ def test_sentinel_ti_provider(ti_lookup):
     res_df = azs_result["RawResult"]
     check.is_instance(res_df, Dict)
     check.is_instance(azs_result["Reference"], str)
-    check.is_true("ThreatIntelligenceIndicator  | where" in azs_result["Reference"])
+    check.is_true("ThreatIntelIndicators  | where" in azs_result["Reference"])
 
     # IP Lookups
     result = ti_lookup.lookup_ioc(ioc=_IOC_IP, start=start, end=end)