contentctl 5.0.0a2__py3-none-any.whl → 5.0.0a3__py3-none-any.whl

contentctl/enrichments/attack_enrichment.py

@@ -1,121 +1,161 @@
-
 from __future__ import annotations
-import sys
 from attackcti import attack_client
 import logging
 from pydantic import BaseModel
 from dataclasses import field
 from typing import Any
 from pathlib import Path
-from contentctl.objects.mitre_attack_enrichment import MitreAttackEnrichment, MitreTactics
+from contentctl.objects.mitre_attack_enrichment import (
+    MitreAttackEnrichment,
+    MitreTactics,
+)
 from contentctl.objects.config import validate
 from contentctl.objects.annotated_types import MITRE_ATTACK_ID_TYPE
-logging.getLogger('taxii2client').setLevel(logging.CRITICAL)
+
+logging.getLogger("taxii2client").setLevel(logging.CRITICAL)
 
 
 class AttackEnrichment(BaseModel):
     data: dict[str, MitreAttackEnrichment] = field(default_factory=dict)
-    use_enrichment:bool = True
-    
+    use_enrichment: bool = True
+
     @staticmethod
-    def getAttackEnrichment(config:validate)->AttackEnrichment:
+    def getAttackEnrichment(config: validate) -> AttackEnrichment:
         enrichment = AttackEnrichment(use_enrichment=config.enrichments)
         _ = enrichment.get_attack_lookup(config.mitre_cti_repo_path, config.enrichments)
         return enrichment
-    
-    def getEnrichmentByMitreID(self, mitre_id:MITRE_ATTACK_ID_TYPE)->MitreAttackEnrichment:
+
+    def getEnrichmentByMitreID(
+        self, mitre_id: MITRE_ATTACK_ID_TYPE
+    ) -> MitreAttackEnrichment:
         if not self.use_enrichment:
-            raise Exception("Error, trying to add Mitre Enrichment, but use_enrichment was set to False")
-        
+            raise Exception(
+                "Error, trying to add Mitre Enrichment, but use_enrichment was set to False"
+            )
+
         enrichment = self.data.get(mitre_id, None)
         if enrichment is not None:
             return enrichment
         else:
-            raise Exception(f"Error, Unable to find Mitre Enrichment for MitreID {mitre_id}")
-        
-    def addMitreIDViaGroupNames(self, technique:dict[str,Any], tactics:list[str], groupNames:list[str])->None:
-        technique_id = technique['technique_id']
-        technique_obj = technique['technique']
+            raise Exception(
+                f"Error, Unable to find Mitre Enrichment for MitreID {mitre_id}"
+            )
+
+    def addMitreIDViaGroupNames(
+        self, technique: dict[str, Any], tactics: list[str], groupNames: list[str]
+    ) -> None:
+        technique_id = technique["technique_id"]
+        technique_obj = technique["technique"]
         tactics.sort()
-        
+
         if technique_id in self.data:
             raise Exception(f"Error, trying to redefine MITRE ID '{technique_id}'")
-        self.data[technique_id] = MitreAttackEnrichment.model_validate({'mitre_attack_id':technique_id, 
-                                                                        'mitre_attack_technique':technique_obj, 
-                                                                        'mitre_attack_tactics':tactics, 
-                                                                        'mitre_attack_groups':groupNames,
-                                                                        'mitre_attack_group_objects':[]}) 
-
-    def addMitreIDViaGroupObjects(self, technique:dict[str,Any], tactics:list[MitreTactics],  groupDicts:list[dict[str,Any]])->None:
-        technique_id = technique['technique_id']
-        technique_obj = technique['technique']
+        self.data[technique_id] = MitreAttackEnrichment.model_validate(
+            {
+                "mitre_attack_id": technique_id,
+                "mitre_attack_technique": technique_obj,
+                "mitre_attack_tactics": tactics,
+                "mitre_attack_groups": groupNames,
+                "mitre_attack_group_objects": [],
+            }
+        )
+
+    def addMitreIDViaGroupObjects(
+        self,
+        technique: dict[str, Any],
+        tactics: list[MitreTactics],
+        groupDicts: list[dict[str, Any]],
+    ) -> None:
+        technique_id = technique["technique_id"]
+        technique_obj = technique["technique"]
         tactics.sort()
-        
-        groupNames:list[str] = sorted([group['group'] for group in groupDicts])
-        
+
+        groupNames: list[str] = sorted([group["group"] for group in groupDicts])
+
         if technique_id in self.data:
             raise Exception(f"Error, trying to redefine MITRE ID '{technique_id}'")
-        
-        self.data[technique_id] = MitreAttackEnrichment.model_validate({'mitre_attack_id': technique_id, 
-                                                                        'mitre_attack_technique': technique_obj, 
-                                                                        'mitre_attack_tactics': tactics, 
-                                                                        'mitre_attack_groups': groupNames,
-                                                                        'mitre_attack_group_objects': groupDicts})
-
-    
-    def get_attack_lookup(self, input_path: Path, enrichments:bool = False) -> dict[str,MitreAttackEnrichment]:            
-        attack_lookup:dict[str,MitreAttackEnrichment] = {}
+
+        self.data[technique_id] = MitreAttackEnrichment.model_validate(
+            {
+                "mitre_attack_id": technique_id,
+                "mitre_attack_technique": technique_obj,
+                "mitre_attack_tactics": tactics,
+                "mitre_attack_groups": groupNames,
+                "mitre_attack_group_objects": groupDicts,
+            }
+        )
+
+    def get_attack_lookup(
+        self, input_path: Path, enrichments: bool = False
+    ) -> dict[str, MitreAttackEnrichment]:
+        attack_lookup: dict[str, MitreAttackEnrichment] = {}
         if not enrichments:
             return attack_lookup
-        
+
         try:
-            print(f"Performing MITRE Enrichment using the repository at {input_path}...",end="", flush=True)
-            # The existence of the input_path is validated during cli argument validation, but it is 
+            print(
+                f"Performing MITRE Enrichment using the repository at {input_path}...",
+                end="",
+                flush=True,
+            )
+            # The existence of the input_path is validated during cli argument validation, but it is
             # possible that the repo is in the wrong format. If the following directories do not
-            # exist, then attack_client will fall back to resolving via REST API. We do not 
-            # want this as it is slow and error prone, so we will force an exception to 
+            # exist, then attack_client will fall back to resolving via REST API. We do not
+            # want this as it is slow and error prone, so we will force an exception to
             # be generated.
-            enterprise_path = input_path/"enterprise-attack"
-            mobile_path = input_path/"ics-attack"
-            ics_path = input_path/"mobile-attack"
-            if not (enterprise_path.is_dir() and mobile_path.is_dir() and ics_path.is_dir()):
-                raise FileNotFoundError("One or more of the following paths does not exist: "
-                                        f"{[str(enterprise_path),str(mobile_path),str(ics_path)]}. "
-                                        f"Please ensure that the {input_path} directory "
-                                        "has been git cloned correctly.") 
+            enterprise_path = input_path / "enterprise-attack"
+            mobile_path = input_path / "ics-attack"
+            ics_path = input_path / "mobile-attack"
+            if not (
+                enterprise_path.is_dir() and mobile_path.is_dir() and ics_path.is_dir()
+            ):
+                raise FileNotFoundError(
+                    "One or more of the following paths does not exist: "
+                    f"{[str(enterprise_path), str(mobile_path), str(ics_path)]}. "
+                    f"Please ensure that the {input_path} directory "
+                    "has been git cloned correctly."
+                )
             lift = attack_client(
-                local_paths= {
-                    "enterprise":str(enterprise_path),
-                    "mobile":str(mobile_path),
-                    "ics":str(ics_path)
+                local_paths={
+                    "enterprise": str(enterprise_path),
+                    "mobile": str(mobile_path),
+                    "ics": str(ics_path),
                 }
             )
-            
-            all_enterprise_techniques = lift.get_enterprise_techniques(stix_format=False)
-            enterprise_relationships = lift.get_enterprise_relationships(stix_format=False)
+
+            all_enterprise_techniques = lift.get_enterprise_techniques(
+                stix_format=False
+            )
+            enterprise_relationships = lift.get_enterprise_relationships(
+                stix_format=False
+            )
             enterprise_groups = lift.get_enterprise_groups(stix_format=False)
-            
+
             for technique in all_enterprise_techniques:
-                apt_groups:list[dict[str,Any]] = []
+                apt_groups: list[dict[str, Any]] = []
                 for relationship in enterprise_relationships:
-                    if (relationship['target_object'] == technique['id']) and relationship['source_object'].startswith('intrusion-set'):
+                    if (
+                        relationship["target_object"] == technique["id"]
+                    ) and relationship["source_object"].startswith("intrusion-set"):
                         for group in enterprise_groups:
-                            if relationship['source_object'] == group['id']:
+                            if relationship["source_object"] == group["id"]:
                                 apt_groups.append(group)
-                                #apt_groups.append(group['group'])
+                                # apt_groups.append(group['group'])
 
                 tactics = []
-                if ('tactic' in technique):
-                    for tactic in technique['tactic']:
-                        tactics.append(tactic.replace('-',' ').title())
+                if "tactic" in technique:
+                    for tactic in technique["tactic"]:
+                        tactics.append(tactic.replace("-", " ").title())
 
                 self.addMitreIDViaGroupObjects(technique, tactics, apt_groups)
-                attack_lookup[technique['technique_id']] = {'technique': technique['technique'], 'tactics': tactics, 'groups': apt_groups}
-
+                attack_lookup[technique["technique_id"]] = {
+                    "technique": technique["technique"],
+                    "tactics": tactics,
+                    "groups": apt_groups,
+                }
 
-        
         except Exception as err:
             raise Exception(f"Error getting MITRE Enrichment: {str(err)}")
-        
+
         print("Done!")
-        return attack_lookup
+        return attack_lookup

contentctl/enrichments/cve_enrichment.py

@@ -1,64 +1,70 @@
 from __future__ import annotations
 from pycvesearch import CVESearch
-import functools
-import os
-import shelve
-import time
-from typing import Annotated, Any, Union, TYPE_CHECKING
-from pydantic import ConfigDict, BaseModel,Field, computed_field
+from typing import Annotated, Union, TYPE_CHECKING
+from pydantic import ConfigDict, BaseModel, Field, computed_field
 from decimal import Decimal
-from requests.exceptions import ReadTimeout
 from contentctl.objects.annotated_types import CVE_TYPE
+
 if TYPE_CHECKING:
     from contentctl.objects.config import validate
 
 
-
-CVESSEARCH_API_URL = 'https://cve.circl.lu'
+CVESSEARCH_API_URL = "https://cve.circl.lu"
 
 
 class CveEnrichmentObj(BaseModel):
     id: CVE_TYPE
-    cvss: Annotated[Decimal, Field(ge=.1, le=10, decimal_places=1)]
+    cvss: Annotated[Decimal, Field(ge=0.1, le=10, decimal_places=1)]
     summary: str
-    
+
     @computed_field
     @property
-    def url(self)->str:
+    def url(self) -> str:
         BASE_NVD_URL = "https://nvd.nist.gov/vuln/detail/"
         return f"{BASE_NVD_URL}{self.id}"
 
 
 class CveEnrichment(BaseModel):
     use_enrichment: bool = True
-    cve_api_obj: Union[CVESearch,None] = None
+    cve_api_obj: Union[CVESearch, None] = None
 
     # Arbitrary_types are allowed to let us use the CVESearch Object
-    model_config = ConfigDict(
-        arbitrary_types_allowed=True,
-        frozen=True
-    )
+    model_config = ConfigDict(arbitrary_types_allowed=True, frozen=True)
 
     @staticmethod
-    def getCveEnrichment(config:validate, timeout_seconds:int=10, force_disable_enrichment:bool=True)->CveEnrichment:
+    def getCveEnrichment(
+        config: validate,
+        timeout_seconds: int = 10,
+        force_disable_enrichment: bool = True,
+    ) -> CveEnrichment:
         if force_disable_enrichment:
-            return CveEnrichment(use_enrichment=False, cve_api_obj=None)    
-        
+            return CveEnrichment(use_enrichment=False, cve_api_obj=None)
+
         if config.enrichments:
             try:
                 cve_api_obj = CVESearch(CVESSEARCH_API_URL, timeout=timeout_seconds)
                 return CveEnrichment(use_enrichment=True, cve_api_obj=cve_api_obj)
             except Exception as e:
-                raise Exception(f"Error setting CVE_SEARCH API to: {CVESSEARCH_API_URL}: {str(e)}")
-        
-        return CveEnrichment(use_enrichment=False, cve_api_obj=None)
-
+                raise Exception(
+                    f"Error setting CVE_SEARCH API to: {CVESSEARCH_API_URL}: {str(e)}"
+                )
 
-    def enrich_cve(self, cve_id:str, raise_exception_on_failure:bool=True)->CveEnrichmentObj:
+        return CveEnrichment(use_enrichment=False, cve_api_obj=None)
 
+    def enrich_cve(
+        self, cve_id: str, raise_exception_on_failure: bool = True
+    ) -> CveEnrichmentObj:
         if not self.use_enrichment:
-            return CveEnrichmentObj(id=cve_id,cvss=Decimal(5.0),summary="SUMMARY NOT AVAILABLE! ONLY THE LINK WILL BE USED AT THIS TIME")
+            return CveEnrichmentObj(
+                id=cve_id,
+                cvss=Decimal(5.0),
+                summary="SUMMARY NOT AVAILABLE! ONLY THE LINK WILL BE USED AT THIS TIME",
+            )
         else:
             print("WARNING - Dynamic enrichment not supported at this time.")
-            return CveEnrichmentObj(id=cve_id,cvss=Decimal(5.0),summary="SUMMARY NOT AVAILABLE! ONLY THE LINK WILL BE USED AT THIS TIME")
-        # Depending on needs, we may add dynamic enrichment functionality back to the tool
+            return CveEnrichmentObj(
+                id=cve_id,
+                cvss=Decimal(5.0),
+                summary="SUMMARY NOT AVAILABLE! ONLY THE LINK WILL BE USED AT THIS TIME",
+            )
+        # Depending on needs, we may add dynamic enrichment functionality back to the tool

contentctl/enrichments/splunk_app_enrichment.py

@@ -1,11 +1,8 @@
 import requests
 import xmltodict
-import json
 import functools
-import pickle
 import shelve
 import os
-import time
 
 SPLUNKBASE_API_URL = "https://apps.splunk.com/api/apps/entriesbyid/"
 
@@ -13,15 +10,16 @@ APP_ENRICHMENT_CACHE_FILENAME = "lookups/APP_ENRICHMENT_CACHE.db"
 
 NON_PERSISTENT_CACHE = {}
 
+
 @functools.cache
-def requests_get_helper(url:str, force_cached_or_offline:bool = False)->bytes:
+def requests_get_helper(url: str, force_cached_or_offline: bool = False) -> bytes:
     if force_cached_or_offline:
         if not os.path.exists(APP_ENRICHMENT_CACHE_FILENAME):
             print(f"Cache at {APP_ENRICHMENT_CACHE_FILENAME} not found - Creating it.")
-        cache = shelve.open(APP_ENRICHMENT_CACHE_FILENAME, flag='c', writeback=True)
+        cache = shelve.open(APP_ENRICHMENT_CACHE_FILENAME, flag="c", writeback=True)
     else:
         cache = NON_PERSISTENT_CACHE
-    
+
     if url in cache:
         req_content = cache[url]
     else:
@@ -29,62 +27,66 @@ def requests_get_helper(url:str, force_cached_or_offline:bool = False)->bytes:
             req = requests.get(url)
             req_content = req.content
             cache[url] = req_content
-        except Exception as e:
-            raise(Exception(f"ERROR - Failed to get Splunk App Enrichment at {SPLUNKBASE_API_URL}"))
-    
+        except Exception:
+            raise (
+                Exception(
+                    f"ERROR - Failed to get Splunk App Enrichment at {SPLUNKBASE_API_URL}"
+                )
+            )
+
     if isinstance(cache, shelve.Shelf):
-        #close the cache if it is a shelf
+        # close the cache if it is a shelf
         cache.close()
-    
+
     return req_content
 
-class SplunkAppEnrichment():
 
+class SplunkAppEnrichment:
     @classmethod
-    def enrich_splunk_app(self, splunk_ta: str, force_cached_or_offline: bool = False) -> dict:
-        
+    def enrich_splunk_app(
+        self, splunk_ta: str, force_cached_or_offline: bool = False
+    ) -> dict:
         appurl = SPLUNKBASE_API_URL + splunk_ta
         splunk_app_enriched = dict()
-        
+
         try:
-            
             content = requests_get_helper(appurl, force_cached_or_offline)
             response_dict = xmltodict.parse(content)
-            
+
             # check if list since data changes depending on answer
             url, results = self._parse_splunkbase_response(response_dict)
             # grab the app name
             for i in results:
-                if i['@name'] == 'appName':
-                    splunk_app_enriched['name'] = i['#text']
-            # grab out the splunkbase url  
-            if 'entriesbyid' in url:
+                if i["@name"] == "appName":
+                    splunk_app_enriched["name"] = i["#text"]
+            # grab out the splunkbase url
+            if "entriesbyid" in url:
                 content = requests_get_helper(url, force_cached_or_offline)
                 response_dict = xmltodict.parse(content)
-                
-                #print(json.dumps(response_dict, indent=2))
+
+                # print(json.dumps(response_dict, indent=2))
                 url, results = self._parse_splunkbase_response(response_dict)
                 # chop the url so we grab the splunkbase portion but not direct download
-                splunk_app_enriched['url'] = url.rsplit('/', 4)[0]
+                splunk_app_enriched["url"] = url.rsplit("/", 4)[0]
         except requests.exceptions.ConnectionError as connErr:
             print(f"There was a connErr for ta {splunk_ta}: {connErr}")
             # there was a connection error lets just capture the name
-            splunk_app_enriched['name'] = splunk_ta
-            splunk_app_enriched['url'] = ''
+            splunk_app_enriched["name"] = splunk_ta
+            splunk_app_enriched["url"] = ""
         except Exception as e:
-            print(f"There was an unknown error enriching the Splunk TA [{splunk_ta}]: {str(e)}")
-            splunk_app_enriched['name'] = splunk_ta
-            splunk_app_enriched['url'] = ''
-
+            print(
+                f"There was an unknown error enriching the Splunk TA [{splunk_ta}]: {str(e)}"
+            )
+            splunk_app_enriched["name"] = splunk_ta
+            splunk_app_enriched["url"] = ""
 
         return splunk_app_enriched
 
     def _parse_splunkbase_response(response_dict):
-        if isinstance(response_dict['feed']['entry'], list):
-            url = response_dict['feed']['entry'][0]['link']['@href']
-            results = response_dict['feed']['entry'][0]['content']['s:dict']['s:key']
+        if isinstance(response_dict["feed"]["entry"], list):
+            url = response_dict["feed"]["entry"][0]["link"]["@href"]
+            results = response_dict["feed"]["entry"][0]["content"]["s:dict"]["s:key"]
         else:
-            url = response_dict['feed']['entry']['link']['@href']
-            results = response_dict['feed']['entry']['content']['s:dict']['s:key']
+            url = response_dict["feed"]["entry"]["link"]["@href"]
+            results = response_dict["feed"]["entry"]["content"]["s:dict"]["s:key"]
         return url, results
-