APP-4521

release_notes.md

@@ -2,11 +2,11 @@
 
 ## 4.0.6
 
--   APP-4522 - [API] Updated batch module to support external date fields
--   APP-4521 - [API] Updated support for IR endpoints
+-   APP-4472 - [API] Added NAICS industry classification module
+-   APP-4482 - [API] Updated transforms to support new static methods
 -   APP-4520 - [Util] Fixed TCEX string operations trim method
--   APP-4482 - [API] Updated transforms to support new static methods.
--   APP-4472 - [NAICS] Added NAICS module for lookup by id or name.
+-   APP-4521 - [API] Updated support for IR endpoints
+-   APP-4522 - [API] Updated batch module to support external date fields
 
 ## 4.0.5
 

tcex/api/tc/v3/_gen/_gen_abc.py

@@ -488,11 +488,7 @@ def tap(self, type_: str):
         ]:
             return 'tcex.api.tc.v3.security'
 
-        if type_.plural().lower() in [
-            'categories',
-            'results',
-            'subtypes',
-        ]:
+        if type_.plural().lower() in ['categories', 'results', 'subtypes', 'keyword_sections']:
             return 'tcex.api.tc.v3.intel_requirements'
 
         return 'tcex.api.tc.v3'

tcex/api/tc/v3/_gen/_gen_object_abc.py

@@ -445,6 +445,53 @@ def as_entity(self) -> dict:
         )
         return '\n'.join(as_entity_property_method)
 
+    def _gen_code_object_replace_type_method(
+        self, type_: str, model_type: str | None = None
+    ) -> str:
+        """Return the method code.
+
+        def replace_artifact(self, **kwargs):
+            '''Replace an Artifact on the object. (mark as staged)
+        """
+        type_ = self.util.snake_string(type_)
+        model_type = self.util.snake_string(model_type or type_)
+        model_reference = model_type
+
+        # get model from map and update requirements
+        model_import_data = self._module_import_data(type_)
+        model_class = model_import_data.get('model_class')
+        self.requirements['first-party'].append(
+            f'''from {model_import_data.get('model_module')} '''
+            f'''import {model_import_data.get('model_class')}'''
+        )
+        stage_method = [
+            (
+                f'''{self.i1}def replace_{model_type.singular()}(self, '''
+                f'''data: dict | list | ObjectABC | {model_class}'''
+                f'''):'''
+            ),
+            f'''{self.i2}"""Replace {type_.singular()} on the object."""''',
+            f'''{self.i2}if not isinstance(data, list):''',
+            f'''{self.i3}data = [data]''',
+            '',
+            f'''{self.i2}if all(isinstance(item, ({model_class}, ObjectABC)) for item in data):'''
+            f'''{self.i3}transformed_data = data''',
+            f'''{self.i2}elif all(isinstance(item, dict) for item in data):'''
+            f'''{self.i3}transformed_data = [{model_class}(**d) for d in data]''',
+            f'''{self.i2}else:'''
+            f'''{self.i3}raise ValueError('Invalid data to replace_{model_type.singular()}')''',
+            '',
+            '',
+            f'''{self.i2}for item in transformed_data:''',
+            f'''{self.i3}item._staged = True''',
+            '',
+            f'''{self.i2}self.model.{model_reference} = transformed_data  # type: ignore''',
+            '',
+            '',
+        ]
+
+        return '\n'.join(stage_method)
+
     def _gen_code_object_stage_type_method(self, type_: str, model_type: str | None = None) -> str:
         """Return the method code.
 
@@ -978,10 +1025,12 @@ def filter ...
             _code += self._gen_code_object_stage_type_method('victim_assets')
 
         # generate stage_associated_group method
-        if 'associatedCases' in add_properties:
+        # ESUP-2532 - Associations are not Bi-Directional for IRs
+        if 'associatedCases' in add_properties and self.type_ != 'intel_requirements':
             _code += self._gen_code_object_stage_type_method('cases', 'associated_cases')
 
-        if 'associatedArtifacts' in add_properties:
+        # ESUP-2532 - Associations are not Bi-Directional for IRs
+        if 'associatedArtifacts' in add_properties and self.type_ != 'intel_requirements':
             _code += self._gen_code_object_stage_type_method('artifacts', 'associated_artifacts')
 
         # victims have associatedGroups but groups must be
@@ -1015,6 +1064,10 @@ def filter ...
         if 'fileOccurrences' in add_properties:
             _code += self._gen_code_object_stage_type_method('file_occurrences')
 
+        # generate stage_keyword_section method
+        if 'keywordSections' in add_properties and self.type_ == 'intel_requirements':
+            _code += self._gen_code_object_replace_type_method('keyword_sections')
+
         # generate stage_note method
         if 'notes' in add_properties:
             _code += self._gen_code_object_stage_type_method('notes')

tcex/api/tc/v3/_gen/model/_property_model.py

@@ -307,7 +307,7 @@ def __process_special_types(cls, pm: 'PropertyModel', extra: dict[str, str]):
                 }
             )
         elif pm.type == 'KeywordSection':
-            bi += 'intel_requirements.keyword_section_model'
+            bi += 'intel_requirements.keyword_sections.keyword_section_model'
             extra.update(
                 {
                     'import_data': f'{bi} import KeywordSectionModel',

tcex/api/tc/v3/intel_requirements/intel_requirement.py

@@ -6,15 +6,16 @@
 
 # first-party
 from tcex.api.tc.v3.api_endpoints import ApiEndpoints
-from tcex.api.tc.v3.artifacts.artifact_model import ArtifactModel
-from tcex.api.tc.v3.cases.case_model import CaseModel
 from tcex.api.tc.v3.groups.group_model import GroupModel
 from tcex.api.tc.v3.indicators.indicator_model import IndicatorModel
 from tcex.api.tc.v3.intel_requirements.intel_requirement_filter import IntelRequirementFilter
 from tcex.api.tc.v3.intel_requirements.intel_requirement_model import (
     IntelRequirementModel,
     IntelRequirementsModel,
 )
+from tcex.api.tc.v3.intel_requirements.keyword_sections.keyword_section_model import (
+    KeywordSectionModel,
+)
 from tcex.api.tc.v3.object_abc import ObjectABC
 from tcex.api.tc.v3.object_collection_abc import ObjectCollectionABC
 from tcex.api.tc.v3.tags.tag_model import TagModel
@@ -140,30 +141,6 @@ def tags(self) -> Generator['Tag', None, None]:
 
         yield from self._iterate_over_sublist(Tags)  # type: ignore
 
-    def stage_associated_case(self, data: dict | ObjectABC | CaseModel):
-        """Stage case on the object."""
-        if isinstance(data, ObjectABC):
-            data = data.model  # type: ignore
-        elif isinstance(data, dict):
-            data = CaseModel(**data)
-
-        if not isinstance(data, CaseModel):
-            raise RuntimeError('Invalid type passed in to stage_associated_case')
-        data._staged = True
-        self.model.associated_cases.data.append(data)  # type: ignore
-
-    def stage_associated_artifact(self, data: dict | ObjectABC | ArtifactModel):
-        """Stage artifact on the object."""
-        if isinstance(data, ObjectABC):
-            data = data.model  # type: ignore
-        elif isinstance(data, dict):
-            data = ArtifactModel(**data)
-
-        if not isinstance(data, ArtifactModel):
-            raise RuntimeError('Invalid type passed in to stage_associated_artifact')
-        data._staged = True
-        self.model.associated_artifacts.data.append(data)  # type: ignore
-
     def stage_associated_group(self, data: dict | ObjectABC | GroupModel):
         """Stage group on the object."""
         if isinstance(data, ObjectABC):
@@ -200,6 +177,23 @@ def stage_associated_indicator(self, data: dict | ObjectABC | IndicatorModel):
         data._staged = True
         self.model.associated_indicators.data.append(data)  # type: ignore
 
+    def replace_keyword_section(self, data: dict | list | ObjectABC | KeywordSectionModel):
+        """Replace keyword_section on the object."""
+        if not isinstance(data, list):
+            data = [data]
+
+        if all(isinstance(item, (KeywordSectionModel, ObjectABC)) for item in data):
+            transformed_data = data
+        elif all(isinstance(item, dict) for item in data):
+            transformed_data = [KeywordSectionModel(**d) for d in data]
+        else:
+            raise ValueError('Invalid data to replace_keyword_section')
+
+        for item in transformed_data:
+            item._staged = True
+
+        self.model.keyword_sections = transformed_data  # type: ignore
+
     def stage_tag(self, data: dict | ObjectABC | TagModel):
         """Stage tag on the object."""
         if isinstance(data, ObjectABC):