refactor: Improve path handling and error messages in ShadowTrackrConnector

- Replaced string manipulation with `Path` for constructing the config file path.
- Enhanced error messages for label creation failures and IP processing.
- Simplified score adjustment logic for false positive estimates.
- Streamlined label handling and description updates for STIX entities.
- Improved exception handling for IP validation.

Author: jabesq

internal-enrichment/shadowtrackr/src/shadowtrackr.py

@@ -2,7 +2,7 @@
 import json
 import os
 from datetime import datetime, timedelta
-from typing import Dict
+from pathlib import Path
 
 import requests
 import yaml
@@ -18,7 +18,7 @@ class ShadowTrackrConnector:
     def __init__(self):
         # Instantiate the connector helper from config
 
-        config_file_path = os.path.dirname(os.path.abspath(__file__)) + "/config.yml"
+        config_file_path = str((Path(__file__).resolve().parent / "config.yml"))
         config = (
             yaml.load(open(config_file_path), Loader=yaml.FullLoader)
             if os.path.isfile(config_file_path)
@@ -71,17 +71,19 @@ def __init__(self):
         )
         if self.label_bogon is None:
             raise ValueError(
-                "The Bogon label could not be created. If your connector does not have the permission to create labels,"
-                " please create it manually before launching"
+                "The Bogon label could not be created. If your connector does not have the permission to create labels, "
+                "please create it manually before launching"
             )
+
         self.label_cloud = self.helper.api.label.read_or_create_unchecked(
             value="cloud", color="#145578"
         )
         if self.label_cloud is None:
             raise ValueError(
-                "The cloud label could not be created. If your connector does not have the permission to create labels,"
-                " please create it manually before launching"
+                "The cloud label could not be created. If your connector does not have the permission to create labels, "
+                "please create it manually before launching."
             )
+
         self.label_cdn = self.helper.api.label.read_or_create_unchecked(
             value="cdn", color="#145578"
         )
@@ -90,6 +92,7 @@ def __init__(self):
                 "The cdn label could not be created. If your connector does not have the permission to create labels, "
                 "please create it manually before launching"
             )
+
         self.label_vpn = self.helper.api.label.read_or_create_unchecked(
             value="vpn", color="#145578"
         )
@@ -98,6 +101,7 @@ def __init__(self):
                 "The vpn label could not be created. If your connector does not have the permission to create labels, "
                 "please create it manually before launching"
             )
+
         self.label_tor = self.helper.api.label.read_or_create_unchecked(
             value="tor", color="#145578"
         )
@@ -119,33 +123,33 @@ def __init__(self):
     def _process_entity(self, stix_objects, stix_entity, opencti_entity) -> str:
         # Search in ShadowTrackr
         is_indicator = False
-        if opencti_entity["entity_type"] == "Indicator":
+        entity_type = opencti_entity["entity_type"]
+        if entity_type == "Indicator":
+            pattern = opencti_entity["pattern"]
             is_indicator = True
-            if "ipv4-addr" in opencti_entity["pattern"]:
-                ip = opencti_entity["pattern"].split("'")[1]
-            elif "ipv6-addr" in opencti_entity["pattern"]:
-                ip = opencti_entity["pattern"].split("'")[1]
+
+            if "ipv4-addr" in pattern or "ipv6-addr" in pattern:
+                ip = pattern.split("'")[1]
             else:
-                return (
-                    "No ip address in indicator, skipping " + opencti_entity["pattern"]
-                )
+                return f"No ip address in indicator, skipping {pattern} {entity_type}"
+
         else:
             ip = opencti_entity["observable_value"]
 
         if self._valid_ip(ip):
             score = opencti_entity["x_opencti_score"]
             old_score = score
-            labels = [l["value"] for l in opencti_entity["objectLabel"]]
+            labels = [label["value"] for label in opencti_entity["objectLabel"]]
             markings = [
-                l["definition"]
-                for l in opencti_entity["objectMarking"]
-                if l["definition_type"] == "TLP"
+                marking["definition"]
+                for marking in opencti_entity["objectMarking"]
+                if marking["definition_type"] == "TLP"
             ]
 
             for tlp in markings:
                 if not OpenCTIConnectorHelper.check_max_tlp(tlp, self.max_tlp):
                     raise ValueError(
-                        "Do not send any data, TLP of the observable is greater than MAX TLP"
+                        f"Do not send any data, TLP of the observable is greater than MAX TLP: {tlp} > {self.max_tlp}"
                     )
 
             if is_indicator:
@@ -155,28 +159,30 @@ def _process_entity(self, stix_objects, stix_entity, opencti_entity) -> str:
 
             if description is not None and "[ShadowTrackr] " in description:
                 return (
-                    "This ip is already processed by the ShadowTrackr connector. We're not doing it again,"
-                    " that might mess up the score"
+                    "This ip is already processed by the ShadowTrackr connector. We're not doing it again, "
+                    "that might mess up the score."
                 )
 
             data = self._check_ip_in_shadowtrackr(ip, labels)
-            if "error" in data:
-                raise Exception("Error: [ShadowTrackr] " + data["error"])
+            if error := data.get("error"):
+                raise ValueError(f"Error: [ShadowTrackr] {error}")
 
             score_lowered = False
             if self.replace_with_lower_score:
-                if data["false_positive_estimate"] > 99:
-                    score -= 60
-                    score_lowered = True
-                elif data["false_positive_estimate"] > 89:
-                    score -= 40
-                    score_lowered = True
-                elif data["false_positive_estimate"] > 69:
-                    score -= 20
-                    score_lowered = True
-                elif data["false_positive_estimate"] > 50:
-                    score -= 10
-                    score_lowered = True
+                false_positive_estimate = data["false_positive_estimate"]
+
+                score_steps = [
+                    (99, 60),
+                    (89, 40),
+                    (69, 20),
+                    (50, 10),
+                ]
+                for threshold, decrement in score_steps:
+                    if false_positive_estimate > threshold:
+                        score -= decrement
+                        score_lowered = True
+                        break
+
                 # set a lower boundary
                 if score < 10:
                     score = 10
@@ -196,127 +202,70 @@ def _process_entity(self, stix_objects, stix_entity, opencti_entity) -> str:
                     timespec="milliseconds"
                 ) + "Z"
                 OpenCTIStix2.put_attribute_in_extension(
-                    stix_entity,
-                    STIX_EXT_OCTI_SCO,
-                    "valid_until",
-                    valid_until,
-                    True,
+                    stix_entity, STIX_EXT_OCTI_SCO, "valid_until", valid_until, True
                 )
                 date_shortened = True
 
-            if data["vpn"]:
-                if is_indicator:
-                    stix_entity["labels"].append("vpn")
-                else:
-                    OpenCTIStix2.put_attribute_in_extension(
-                        stix_entity,
-                        STIX_EXT_OCTI_SCO,
-                        "labels",
-                        "vpn",
-                        True,
-                    )
-            if data["cdn"]:
-                if is_indicator:
-                    stix_entity["labels"].append("cdn")
-                else:
-                    OpenCTIStix2.put_attribute_in_extension(
-                        stix_entity,
-                        STIX_EXT_OCTI_SCO,
-                        "labels",
-                        "cdn",
-                        True,
-                    )
-            if data["cloud"]:
-                if is_indicator:
-                    stix_entity["labels"].append("cloud")
-                else:
-                    OpenCTIStix2.put_attribute_in_extension(
-                        stix_entity,
-                        STIX_EXT_OCTI_SCO,
-                        "labels",
-                        "cloud",
-                        True,
-                    )
-            if data["bogon"]:
-                if is_indicator:
-                    stix_entity["labels"].append("bogon")
-                else:
-                    OpenCTIStix2.put_attribute_in_extension(
-                        stix_entity,
-                        STIX_EXT_OCTI_SCO,
-                        "labels",
-                        "bogon",
-                        True,
-                    )
-            if data["tor"]:
-                if is_indicator:
-                    stix_entity["labels"].append("tor")
-                else:
-                    OpenCTIStix2.put_attribute_in_extension(
-                        stix_entity,
-                        STIX_EXT_OCTI_SCO,
-                        "labels",
-                        "tor",
-                        True,
-                    )
-            if data["public_dns"]:
-                if is_indicator:
-                    stix_entity["labels"].append("public dns server")
-                else:
-                    OpenCTIStix2.put_attribute_in_extension(
-                        stix_entity,
-                        STIX_EXT_OCTI_SCO,
-                        "labels",
-                        "public dns server",
-                        True,
-                    )
+            for label in ["vpn", "cdn", "cloud", "bogon", "tor", "public_dns"]:
+                if data[label]:
+                    if is_indicator:
+                        stix_entity["labels"].append(label)
+                    else:
+                        OpenCTIStix2.put_attribute_in_extension(
+                            stix_entity, STIX_EXT_OCTI_SCO, "labels", label, True
+                        )
 
             text = ""
             if data["cdn"]:
                 if data["cdn_provider"]:
-                    text = (
-                        "This is an ip address in the " + data["cdn_provider"] + " CDN."
-                    )
+                    text = f"This is an ip address in the {data['cdn_provider']} CDN."
                 else:
                     text = "This is an ip address in a CDN."
-                text += " CDN ip addresses change regularly, and are not very useful to track. "
-                if score_lowered and date_shortened:
-                    text += " The score is adjusted downwards, and the valid until date set to 1 day."
-                elif score_lowered:
-                    text += " The score is adjusted downwards."
+
+                text += " CDN ip addresses change regularly, and are not very useful to track."
+                if score_lowered:
+                    if date_shortened:
+                        text += " The score is adjusted downwards, and the valid until date set to 1 day."
+                    else:
+                        text += " The score is adjusted downwards."
 
             elif data["cloud"]:
                 if data["cloud_provider"]:
                     text = (
-                        "This is an ip address in the "
-                        + data["cloud_provider"]
-                        + " cloud."
+                        f" This is an ip address in the {data['cloud_provider']} cloud."
                     )
                 else:
-                    text = "This is an ip address in a cloud."
-                text += " Cloud ip addresses change regularly, and are not very useful to track. "
-                if score_lowered and date_shortened:
-                    text += " The score is adjusted downwards, and the valid until date set to 1 day."
-                elif score_lowered:
-                    text += " The score is adjusted downwards."
+                    text = " This is an ip address in a cloud."
+
+                text += " Cloud ip addresses change regularly, and are not very useful to track."
+                if score_lowered:
+                    if date_shortened:
+                        text += " The score is adjusted downwards, and the valid until date set to 1 day."
+                    else:
+                        text += " The score is adjusted downwards."
 
             elif data["vpn"]:
-                text = "This ip address is a VPN."
-                text += " VPN ip addresses change regularly, and are not very useful to track. "
-                if score_lowered and date_shortened:
-                    text += " The score is adjusted downwards, and the valid until date set to 1 day."
-                elif score_lowered:
-                    text += " The score is adjusted downwards."
+                text = (
+                    "This ip address is a VPN. "
+                    "VPN ip addresses change regularly, and are not very useful to track."
+                )
+                if score_lowered:
+                    if date_shortened:
+                        text += " The score is adjusted downwards, and the valid until date set to 1 day."
+                    else:
+                        text += " The score is adjusted downwards."
 
             elif data["public_dns"]:
-                text = "This ip address is a public DNS server."
-                text += " Public DNS servers are often used in malware to check for an internet connection, "
-                text += "and automated analysis tools regularly extract them as indicators. This is not very useful."
+                text = (
+                    "This ip address is a public DNS server. "
+                    "Public DNS servers are often used in malware to check for an internet connection, "
+                    "and automated analysis tools regularly extract them as indicators. This is not very useful."
+                )
                 if self.replace_with_lower_score:
                     text += " The score is adjusted downwards."
 
             if text:
-                description += "\n[ShadowTrackr] " + text
+                description += f"\n[ShadowTrackr] {text}"
                 if is_indicator:
                     stix_entity["description"] = description
                 else:
@@ -332,39 +281,23 @@ def _process_entity(self, stix_objects, stix_entity, opencti_entity) -> str:
                 if date_shortened:
                     serialized_bundle = self.helper.stix2_create_bundle(stix_objects)
                     self.helper.send_stix2_bundle(serialized_bundle)
-                    return (
-                        "Found data on "
-                        + ip
-                        + ". Score not changed, but valid_until shortened to 1 day."
-                    )
+                    return f"Found data on {ip}. Score not changed, but valid_until shortened to 1 day."
                 else:
-                    return "Found data on " + ip + ". Score not changed."
+                    return f"Found data on {ip}. Score not changed."
             else:
                 serialized_bundle = self.helper.stix2_create_bundle(stix_objects)
                 self.helper.send_stix2_bundle(serialized_bundle)
                 if date_shortened:
                     return (
-                        "Found data on "
-                        + ip
-                        + ". Score changed from "
-                        + str(old_score)
-                        + " to "
-                        + str(score)
-                        + ", valid_until shortened to 1 day."
+                        f"Found data on {ip}. Score changed from {old_score} to {score}, "
+                        "valid_until shortened to 1 day."
                     )
                 else:
-                    return (
-                        "Found data on "
-                        + ip
-                        + ". Score changed from "
-                        + str(old_score)
-                        + " to "
-                        + str(score)
-                    )
+                    return f"Found data on {ip}. Score changed from {old_score} to {score}."
 
-        return "Invalid ip: " + ip
+        return f"Invalid ip: {ip}"
 
-    def _process_message(self, data: Dict) -> str:
+    def _process_message(self, data: dict) -> str:
         stix_objects = data["stix_objects"]
         stix_entity = data["stix_entity"]
         opencti_entity = data["enrichment_entity"]
@@ -381,18 +314,18 @@ def _check_ip_in_shadowtrackr(self, ip, labels=None) -> dict:
             labels = ""
 
         base_url = "https://shadowtrackr.com/api/v3/ip_info"
-        full_url = (
-            base_url + "?api_key=" + self.api_key + "&ip=" + ip + "&tags=" + labels
+        response = requests.get(
+            base_url, params={"api_key": self.api_key, "ip": ip, "tags": labels or ""}
         )
-        response = requests.get(full_url)
         data = json.loads(response.text)
         return data
 
     def _valid_ip(self, ip):
         try:
             ipaddress.ip_address(ip)
             return True
-        except:
+        except Exception as e:
+            self.helper.log_error("Error validating ip", {"error": str(e)})
             return False