contentctl 3.6.0__py3-none-any.whl → 4.0.2__py3-none-any.whl

contentctl/enrichments/attack_enrichment.py

@@ -1,4 +1,5 @@
 
+from __future__ import annotations
 import csv
 import os
 from posixpath import split
@@ -6,16 +7,57 @@ from typing import Optional
 import sys
 from attackcti import attack_client
 import logging
+from pydantic import BaseModel, Field
+from dataclasses import field
+from typing import Union,Annotated
+from contentctl.objects.mitre_attack_enrichment import MitreAttackEnrichment
+from contentctl.objects.config import validate
 logging.getLogger('taxii2client').setLevel(logging.CRITICAL)
 
 
-class AttackEnrichment():
+class AttackEnrichment(BaseModel):
+    data: dict[str, MitreAttackEnrichment] = field(default_factory=dict)
+    use_enrichment:bool = True
+    
+    @staticmethod
+    def getAttackEnrichment(config:validate)->AttackEnrichment:
+        enrichment = AttackEnrichment(use_enrichment=config.enrichments)
+        _ = enrichment.get_attack_lookup(str(config.path))
+        return enrichment
+    
+    def getEnrichmentByMitreID(self, mitre_id:Annotated[str, Field(pattern="^T\d{4}(.\d{3})?$")])->Union[MitreAttackEnrichment,None]:
+        if not self.use_enrichment:
+            return None
+        
+        enrichment = self.data.get(mitre_id, None)
+        if enrichment is not None:
+            return enrichment
+        else:
+            raise ValueError(f"Error, Unable to find Mitre Enrichment for MitreID {mitre_id}")
+        
 
-    @classmethod
-    def get_attack_lookup(self, input_path: str, store_csv = False, force_cached_or_offline: bool = False, skip_enrichment:bool = False) -> dict:
+    def addMitreID(self, technique:dict, tactics:list[str], groups:list[str])->None:
+        
+        technique_id = technique['technique_id']
+        technique_obj = technique['technique']
+        tactics.sort()
+        groups.sort()
+
+        if technique_id in self.data:
+            raise ValueError(f"Error, trying to redefine MITRE ID '{technique_id}'")
+        
+        self.data[technique_id] = MitreAttackEnrichment(mitre_attack_id=technique_id, 
+                                                        mitre_attack_technique=technique_obj, 
+                                                        mitre_attack_tactics=tactics, 
+                                                        mitre_attack_groups=groups)
+
+    
+    def get_attack_lookup(self, input_path: str, store_csv: bool = False, force_cached_or_offline: bool = False, skip_enrichment:bool = False) -> dict:
+        if self.use_enrichment is False:
+            return {}
         print("Getting MITRE Attack Enrichment Data. This may take some time...")
         attack_lookup = dict()
-        file_path = os.path.join(input_path, "lookups", "mitre_enrichment.csv")
+        file_path = os.path.join(input_path, "app_template", "lookups", "mitre_enrichment.csv")
 
         if skip_enrichment is True:
             print("Skipping enrichment")
@@ -28,36 +70,38 @@ class AttackEnrichment():
             lift = attack_client()
             print(f"\r{'Client'.rjust(23)}: [{100:3.0f}%]...Done!", end="\n", flush=True)
             
-            print(f"\r{'Enterprise'.rjust(23)}: [{0.0:3.0f}%]...", end="", flush=True)
-            all_enterprise = lift.get_enterprise(stix_format=False)
-            print(f"\r{'Enterprise'.rjust(23)}: [{100:3.0f}%]...Done!", end="\n", flush=True)
+            print(f"\r{'Techniques'.rjust(23)}: [{0.0:3.0f}%]...", end="", flush=True)
+            all_enterprise_techniques = lift.get_enterprise_techniques(stix_format=False)
+            
+            print(f"\r{'Techniques'.rjust(23)}: [{100:3.0f}%]...Done!", end="\n", flush=True)
             
             print(f"\r{'Relationships'.rjust(23)}: [{0.0:3.0f}%]...", end="", flush=True)
-            enterprise_relationships = lift.get_enterprise_relationships()
+            enterprise_relationships = lift.get_enterprise_relationships(stix_format=False)
             print(f"\r{'Relationships'.rjust(23)}: [{100:3.0f}%]...Done!", end="\n", flush=True)
             
             print(f"\r{'Groups'.rjust(23)}: [{0:3.0f}%]...", end="", flush=True)
-            enterprise_groups = lift.get_enterprise_groups()
+            enterprise_groups = lift.get_enterprise_groups(stix_format=False)
             print(f"\r{'Groups'.rjust(23)}: [{100:3.0f}%]...Done!", end="\n", flush=True)
             
-            for index, technique in enumerate(all_enterprise['techniques']):
-                progress_percent = ((index+1)/len(all_enterprise['techniques'])) * 100
+            
+            for index, technique in enumerate(all_enterprise_techniques):
+                progress_percent = ((index+1)/len(all_enterprise_techniques)) * 100
                 if (sys.stdout.isatty() and sys.stdin.isatty() and sys.stderr.isatty()):
                     print(f"\r\t{'MITRE Technique Progress'.rjust(23)}: [{progress_percent:3.0f}%]...", end="", flush=True)
                 apt_groups = []
                 for relationship in enterprise_relationships:
-                    if (relationship['target_ref'] == technique['id']) and relationship['source_ref'].startswith('intrusion-set'):
+                    if (relationship['target_object'] == technique['id']) and relationship['source_object'].startswith('intrusion-set'):
                         for group in enterprise_groups:
-                            if relationship['source_ref'] == group['id']:
-                                apt_groups.append(group['name'])
+                            if relationship['source_object'] == group['id']:
+                                apt_groups.append(group['group'])
 
                 tactics = []
                 if ('tactic' in technique):
                     for tactic in technique['tactic']:
                         tactics.append(tactic.replace('-',' ').title())
 
-                if not ('revoked' in technique):
-                    attack_lookup[technique['technique_id']] = {'technique': technique['technique'], 'tactics': tactics, 'groups': apt_groups}
+                self.addMitreID(technique, tactics, apt_groups)
+                attack_lookup[technique['technique_id']] = {'technique': technique['technique'], 'tactics': tactics, 'groups': apt_groups}
 
             if store_csv:
                 f = open(file_path, 'w')
@@ -79,13 +123,19 @@ class AttackEnrichment():
                 f.close()
 
         except Exception as err:
-            print('Warning: ' + str(err))
-            print('Use local copy lookups/mitre_enrichment.csv')
-            dict_from_csv = {}
+            print(f'\nError: {str(err)}')
+            print('Use local copy app_template/lookups/mitre_enrichment.csv')
             with open(file_path, mode='r') as inp:
                 reader = csv.reader(inp)
                 attack_lookup = {rows[0]:{'technique': rows[1], 'tactics': rows[2].split('|'), 'groups': rows[3].split('|')} for rows in reader}
             attack_lookup.pop('mitre_id')
+            for key in attack_lookup.keys():
+                technique_input = {'technique_id': key , 'technique': attack_lookup[key]['technique'] }
+                tactics_input = attack_lookup[key]['tactics']
+                groups_input = attack_lookup[key]['groups']
+                self.addMitreID(technique=technique_input, tactics=tactics_input, groups=groups_input)
+            
+                
 
         print("Done!")
         return attack_lookup

contentctl/enrichments/cve_enrichment.py

@@ -1,10 +1,13 @@
-
+from __future__ import annotations
 from pycvesearch import CVESearch
 import functools
 import os
 import shelve
 import time
-import sys
+from typing import Annotated
+from pydantic import BaseModel,Field,ConfigDict
+
+from decimal import Decimal
 CVESSEARCH_API_URL = 'https://cve.circl.lu'
 
 CVE_CACHE_FILENAME = "lookups/CVE_CACHE.db"
@@ -12,7 +15,7 @@ CVE_CACHE_FILENAME = "lookups/CVE_CACHE.db"
 NON_PERSISTENT_CACHE = {}
 
 
-
+''''''
 @functools.cache
 def cvesearch_helper(url:str, cve_id:str, force_cached_or_offline:bool=False, max_api_attempts:int=3, retry_sleep_seconds:int=5):
     if max_api_attempts < 1:
@@ -63,10 +66,22 @@ def cvesearch_id_helper(url:str):
 
 
 
-class CveEnrichment():
 
+class CveEnrichmentObj(BaseModel):
+    id:Annotated[str, "^CVE-[1|2][0-9]{3}-[0-9]+$"]
+    cvss:Annotated[Decimal, Field(ge=.1, le=10, decimal_places=1)]
+    summary:str
+
+
+    @staticmethod
+    def buildEnrichmentOnFailure(id:Annotated[str, "^CVE-[1|2][0-9]{3}-[0-9]+$"], errorMessage:str)->CveEnrichmentObj:
+        message = f"{errorMessage}. Default CVSS of 5.0 used"
+        print(message)
+        return CveEnrichmentObj(id=id, cvss=Decimal(5.0), summary=message)
+
+class CveEnrichment():
     @classmethod
-    def enrich_cve(self, cve_id: str, force_cached_or_offline: bool = False) -> dict:
+    def enrich_cve(cls, cve_id: str, force_cached_or_offline: bool = False, treat_failures_as_warnings:bool=True) -> CveEnrichmentObj:
         cve_enriched = dict()
         try:
             
@@ -74,12 +89,12 @@ class CveEnrichment():
             cve_enriched['id'] = cve_id
             cve_enriched['cvss'] = result['cvss']
             cve_enriched['summary'] = result['summary']
-        except TypeError as TypeErr:
-            # there was a error calling the circl api lets just empty the object
-            print("WARNING, issue enriching {0}, with error: {1}".format(cve_id, str(TypeErr)))
-            cve_enriched = dict()
-            
         except Exception as e:
-            print("WARNING - {0}".format(str(e)))
-    
-        return cve_enriched
+            message = f"issue enriching {cve_id}, with error: {str(e)}"
+            if treat_failures_as_warnings: 
+                return CveEnrichmentObj.buildEnrichmentOnFailure(id = cve_id, errorMessage=f"WARNING, {message}")
+            else:
+                raise ValueError(f"ERROR, {message}")
+        
+        return CveEnrichmentObj.model_validate(cve_enriched)
+        

contentctl/helper/link_validator.py

@@ -1,8 +1,5 @@
-import re
-from tracemalloc import start
-from unittest.mock import DEFAULT
-from pydantic import BaseModel, validator, root_validator,Field
-from typing import Union, Callable
+from pydantic import BaseModel, model_validator
+from typing import Union, Callable, Any
 import requests
 import urllib3, urllib3.exceptions
 import time
@@ -14,6 +11,7 @@ import shelve
 DEFAULT_USER_AGENT_STRING = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/101.0.4951.41 Safari/537.36"
 ALLOWED_HTTP_CODES = [200]
 class LinkStats(BaseModel):
+
     #Static Values
     method: Callable = requests.get
     allowed_http_codes: list[int] = ALLOWED_HTTP_CODES 
@@ -42,17 +40,17 @@ class LinkStats(BaseModel):
         self.referencing_files.add(referencing_file)
         return self.valid
     
-    @root_validator
-    def check_reference(cls, values):
+    @model_validator(mode="before")
+    def check_reference(cls, data:Any)->Any:
         start_time = time.time()
         #Get out all the fields names to make them easier to reference
-        method = values['method']
-        reference = values['reference']
-        timeout_seconds = values['timeout_seconds']
-        headers = values['headers']
-        allow_redirects = values['allow_redirects']
-        verify_ssl = values['verify_ssl']
-        allowed_http_codes = values['allowed_http_codes']
+        method = data['method']
+        reference = data['reference']
+        timeout_seconds = data['timeout_seconds']
+        headers = data['headers']
+        allow_redirects = data['allow_redirects']
+        verify_ssl = data['verify_ssl']
+        allowed_http_codes = data['allowed_http_codes']
         if not (reference.startswith("http://") or reference.startswith("https://")):
             raise(ValueError(f"Reference {reference} does not begin with http(s). Only http(s) references are supported"))
         
@@ -61,29 +59,29 @@ class LinkStats(BaseModel):
                             headers = headers, 
                             allow_redirects=allow_redirects, verify=verify_ssl)
             resolution_time = time.time() - start_time
-            values['status_code'] = get.status_code
-            values['resolution_time'] = resolution_time
+            data['status_code'] = get.status_code
+            data['resolution_time'] = resolution_time
             if reference != get.url:
-                values['redirect'] = get.url
+                data['redirect'] = get.url
             else:
-                values['redirect'] = None #None is also already the default
+                data['redirect'] = None #None is also already the default
 
             #Returns the updated values and sets them for the object
             if get.status_code in allowed_http_codes:
-                values['valid'] = True
+                data['valid'] = True
             else:
                 #print(f"Unacceptable HTTP Status Code {get.status_code} received for {reference}")
-                values['valid'] = False
-            return values    
+                data['valid'] = False
+            return data    
 
         except Exception as e:
             resolution_time = time.time() - start_time
             #print(f"Reference {reference} was not reachable after {resolution_time:.2f} seconds")
-            values['status_code'] = 0
-            values['valid'] = False
-            values['redirect'] = None
-            values['resolution_time'] = resolution_time
-            return values
+            data['status_code'] = 0
+            data['valid'] = False
+            data['redirect'] = None
+            data['resolution_time'] = resolution_time
+            return data
 
 
 class LinkValidator(abc.ABC):

contentctl/helper/utils.py

@@ -6,13 +6,17 @@ import random
 import string
 from timeit import default_timer
 import pathlib
-import datetime
+
 from typing import Union, Tuple
-from pydantic import ValidationError
 import tqdm
-from contentctl.objects.security_content_object import SecurityContentObject
 from math import ceil
 
+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from contentctl.objects.security_content_object import SecurityContentObject
+from contentctl.objects.security_content_object import SecurityContentObject
+
+
 TOTAL_BYTES = 0
 ALWAYS_PULL = True