gis-9415 Add alternative translations functionality

nazargesyk · nazargesyk · commit e4602e017868 · 2025-01-06T15:35:59.000+02:00
diff --git a/uncoder-core/app/routers/translate.py b/uncoder-core/app/routers/translate.py
@@ -3,23 +3,29 @@
 from app.models.translation import InfoMessage, OneTranslationData, Platform, TranslatorPlatforms
 from app.translator.core.context_vars import return_only_first_query_ctx_var
 from app.translator.cti_translator import CTITranslator
-from app.translator.translator import Translator
+from app.translator.translator import app_translator
 
 st_router = APIRouter()
 
-translator = Translator()
-
 
 @st_router.post("/translate", tags=["translator"], description="Generate target translation")
 @st_router.post("/translate/", include_in_schema=False)
 def translate_one(
     source_platform_id: str = Body(..., embed=True),
+    source_alt_mapping: dict = Body(None, embed=True),
     target_platform_id: str = Body(..., embed=True),
+    target_alt_mapping: dict = Body(None, embed=True),
     text: str = Body(..., embed=True),
     return_only_first_query: bool = False,
 ) -> OneTranslationData:
     return_only_first_query_ctx_var.set(return_only_first_query)
-    status, data = translator.translate_one(text=text, source=source_platform_id, target=target_platform_id)
+    status, data = app_translator.translate_one(
+        text=text,
+        source=source_platform_id,
+        target=target_platform_id,
+        source_alt_mapping=source_alt_mapping,
+        target_alt_mapping=target_alt_mapping,
+    )
     if status:
         return OneTranslationData(status=status, translation=data, target_platform_id=target_platform_id)
 
@@ -35,7 +41,7 @@ def translate_all(
     return_only_first_query: bool = False,
 ) -> list[OneTranslationData]:
     return_only_first_query_ctx_var.set(return_only_first_query)
-    result = translator.translate_all(text=text, source=source_platform_id)
+    result = app_translator.translate_all(text=text, source=source_platform_id)
     translations = []
     for platform_result in result:
         if platform_result.get("status"):
@@ -60,14 +66,14 @@ def translate_all(
 @st_router.get("/platforms", tags=["translator"], description="Get translator platforms")
 @st_router.get("/platforms/", include_in_schema=False)
 def get_translator_platforms() -> TranslatorPlatforms:
-    renders, parsers = translator.get_all_platforms()
+    renders, parsers = app_translator.get_all_platforms()
     return TranslatorPlatforms(renders=renders, parsers=parsers)
 
 
 @st_router.get("/all_platforms", description="Get Sigma, RootA and iocs platforms")
 @st_router.get("/all_platforms/", include_in_schema=False)
 def get_all_platforms() -> list:
-    translator_renders, translator_parsers = translator.get_all_platforms()
+    translator_renders, translator_parsers = app_translator.get_all_platforms()
     return [
         Platform(
             id="roota",
diff --git a/uncoder-core/app/translator/core/mapping.py b/uncoder-core/app/translator/core/mapping.py
@@ -112,17 +112,23 @@ class BasePlatformMappings:
 
     def __init__(self, platform_dir: str, platform_details: PlatformDetails):
         self._loader = LoaderFileMappings()
-        self._platform_dir = platform_dir
         self.details = platform_details
-        self._source_mappings = self.prepare_mapping()
+        self._source_mappings = self.prepare_mapping(platform_dir)
+        self._alternative_mappings = self.prepare_alternative_mapping(platform_dir)
 
     def update_default_source_mapping(self, default_mapping: SourceMapping, fields_mapping: FieldsMapping) -> None:
         default_mapping.fields_mapping.update(fields_mapping)
 
-    def prepare_mapping(self) -> dict[str, SourceMapping]:
+    def prepare_alternative_mapping(self, platform_dir: str) -> dict[str, dict[str, SourceMapping]]:
+        alternative_mappings = {}
+        for name, platform_dir in self._loader.get_platform_alternative_mappings(platform_dir).items():
+            alternative_mappings[name] = self.prepare_mapping(platform_dir)
+        return alternative_mappings
+
+    def prepare_mapping(self, platform_dir: str) -> dict[str, SourceMapping]:
         source_mappings = {}
         default_mapping = SourceMapping(source_id=DEFAULT_MAPPING_NAME)
-        for mapping_dict in self._loader.load_platform_mappings(self._platform_dir):
+        for mapping_dict in self._loader.load_platform_mappings(platform_dir):
             log_source_signature = self.prepare_log_source_signature(mapping=mapping_dict)
             if (source_id := mapping_dict["source"]) == DEFAULT_MAPPING_NAME:
                 default_mapping.log_source_signature = log_source_signature
@@ -183,6 +189,9 @@ def get_source_mappings_by_fields_and_log_sources(
     def get_source_mapping(self, source_id: str) -> Optional[SourceMapping]:
         return self._source_mappings.get(source_id)
 
+    def get_alternative_source_mapping(self, alt_config_name: str, source_id: str) -> Optional[SourceMapping]:
+        return self._alternative_mappings.get(alt_config_name, {}).get(source_id)
+
     def get_source_mappings_by_ids(
         self, source_mapping_ids: list[str], return_default: bool = True
     ) -> list[SourceMapping]:
@@ -198,6 +207,21 @@ def get_source_mappings_by_ids(
 
         return source_mappings
 
+    def get_alternative_source_mappings_by_ids(
+        self, source_mapping_ids: list[str], alt_config_name: str, return_default: bool = True
+    ) -> list[SourceMapping]:
+        source_mappings = []
+        for source_mapping_id in source_mapping_ids:
+            if source_mapping_id == DEFAULT_MAPPING_NAME:
+                continue
+            if source_mapping := self.get_alternative_source_mapping(alt_config_name, source_mapping_id):
+                source_mappings.append(source_mapping)
+
+        if not source_mappings and return_default:
+            source_mappings = [self.get_alternative_source_mapping(alt_config_name, DEFAULT_MAPPING_NAME)]
+
+        return source_mappings
+
     def get_source_mappings_by_log_sources(self, log_sources: dict) -> Optional[list[str]]:
         raise NotImplementedError("Abstract method")
 
@@ -249,11 +273,11 @@ def map_field(field: Field, source_mapping: SourceMapping) -> list[str]:
 
 
 class BaseCommonPlatformMappings(ABC, BasePlatformMappings):
-    def prepare_mapping(self) -> dict[str, SourceMapping]:
+    def prepare_mapping(self, platform_dir: str) -> dict[str, SourceMapping]:
         source_mappings = {}
-        common_field_mapping = self._loader.load_common_mapping(self._platform_dir).get("field_mapping", {})
+        common_field_mapping = self._loader.load_common_mapping(platform_dir).get("field_mapping", {})
 
-        for mapping_dict in self._loader.load_platform_mappings(self._platform_dir):
+        for mapping_dict in self._loader.load_platform_mappings(platform_dir):
             source_id = mapping_dict["source"]
             log_source_signature = self.prepare_log_source_signature(mapping=mapping_dict)
             fields_mapping = self.prepare_fields_mapping(field_mapping=common_field_mapping)
diff --git a/uncoder-core/app/translator/core/mixins/tokens.py b/uncoder-core/app/translator/core/mixins/tokens.py
@@ -1,5 +1,3 @@
-from typing import Union
-
 from app.translator.core.const import QUERY_TOKEN_TYPE
 from app.translator.core.custom_types.tokens import LogicalOperatorType, OperatorType
 from app.translator.core.mapping import SourceMapping
diff --git a/uncoder-core/app/translator/core/models/query_container.py b/uncoder-core/app/translator/core/models/query_container.py
@@ -75,6 +75,8 @@ def __init__(
         status: Optional[str] = None,
         false_positives: Optional[list[str]] = None,
         source_mapping_ids: Optional[list[str]] = None,
+        source_alt_mapping: Optional[str] = None,
+        target_alt_mapping: Optional[str] = None,
         parsed_logsources: Optional[dict] = None,
         timeframe: Optional[timedelta] = None,
         query_period: Optional[timedelta] = None,
@@ -107,6 +109,8 @@ def __init__(
         self.timeframe = timeframe
         self.query_period = query_period
         self.raw_metainfo_container = raw_metainfo_container or RawMetaInfoContainer()
+        self.source_alt_mapping = source_alt_mapping
+        self.target_alt_mapping = target_alt_mapping
 
     @property
     def author_str(self) -> str:
diff --git a/uncoder-core/app/translator/core/render.py b/uncoder-core/app/translator/core/render.py
@@ -463,7 +463,15 @@ def _generate_from_tokenized_query_container_by_source_mapping(
     def generate_from_tokenized_query_container(self, query_container: TokenizedQueryContainer) -> str:
         queries_map = {}
         errors = []
-        source_mappings = self.mappings.get_source_mappings_by_ids(query_container.meta_info.source_mapping_ids)
+        if query_container.meta_info.target_alt_mapping:
+            source_mappings = self.mappings.get_alternative_source_mappings_by_ids(
+                source_mapping_ids=query_container.meta_info.source_mapping_ids,
+                alt_config_name=query_container.meta_info.target_alt_mapping,
+            )
+        else:
+            source_mappings = self.mappings.get_source_mappings_by_ids(
+                source_mapping_ids=query_container.meta_info.source_mapping_ids
+            )
 
         for source_mapping in source_mappings:
             try:
diff --git a/uncoder-core/app/translator/mappings/platforms/microsoft_sentinel/alternative/cloudtrail/cloudtrail.yml b/uncoder-core/app/translator/mappings/platforms/microsoft_sentinel/alternative/cloudtrail/cloudtrail.yml
@@ -0,0 +1,13 @@
+platform: Microsoft Sentinel
+source: cloudtrail
+
+
+log_source:
+  table: [AWSCloudTrail]
+
+default_log_source:
+  table: AWSCloudTrail
+
+field_mapping:
+  eventSource: EventSource
+  eventName: EventName
diff --git a/uncoder-core/app/translator/mappings/platforms/microsoft_sentinel/alternative/cloudtrail/default.yml b/uncoder-core/app/translator/mappings/platforms/microsoft_sentinel/alternative/cloudtrail/default.yml
@@ -0,0 +1,7 @@
+platform: Microsoft Sentinel
+source: default
+
+
+
+default_log_source:
+  table: union *
diff --git a/uncoder-core/app/translator/mappings/platforms/microsoft_sentinel/alternative/ossem/default.yml b/uncoder-core/app/translator/mappings/platforms/microsoft_sentinel/alternative/ossem/default.yml
@@ -0,0 +1,7 @@
+platform: Microsoft Sentinel
+source: default
+
+
+
+default_log_source:
+  table: union *
diff --git a/uncoder-core/app/translator/mappings/utils/load_from_files.py b/uncoder-core/app/translator/mappings/utils/load_from_files.py
@@ -7,6 +7,7 @@
 
 COMMON_FIELD_MAPPING_FILE_NAME = "common.yml"
 DEFAULT_FIELD_MAPPING_FILE_NAME = "default.yml"
+DEFAULT_ALTERNATIVE_MAPPINGS_FOLDER_NAME = "alternative"
 
 
 class LoaderFileMappings:
@@ -21,10 +22,22 @@ def load_mapping(mapping_file_path: str) -> dict:
                 print(err)
                 return {}
 
+    def get_platform_alternative_mappings(self, platform_dir: str) -> dict[str:str]:
+        platform_path = os.path.join(self.base_mapping_filepath, platform_dir, DEFAULT_ALTERNATIVE_MAPPINGS_FOLDER_NAME)
+        for folders in os.walk(platform_path):
+            result = {}
+            for folder in folders[1]:
+                result[folder] = os.path.join(platform_dir, DEFAULT_ALTERNATIVE_MAPPINGS_FOLDER_NAME, folder)
+            return result
+        return {}
+
     def load_platform_mappings(self, platform_dir: str) -> Generator[dict, None, None]:
         platform_path = os.path.join(self.base_mapping_filepath, platform_dir)
         for mapping_file in os.listdir(platform_path):
-            if mapping_file not in (COMMON_FIELD_MAPPING_FILE_NAME, DEFAULT_FIELD_MAPPING_FILE_NAME):
+            if mapping_file.endswith(".yml") and mapping_file not in (
+                COMMON_FIELD_MAPPING_FILE_NAME,
+                DEFAULT_FIELD_MAPPING_FILE_NAME,
+            ):
                 yield self.load_mapping(mapping_file_path=os.path.join(platform_path, mapping_file))
         yield self.load_mapping(mapping_file_path=os.path.join(platform_path, DEFAULT_FIELD_MAPPING_FILE_NAME))
 
diff --git a/uncoder-core/app/translator/platforms/arcsight/renders/arcsight_cti.py b/uncoder-core/app/translator/platforms/arcsight/renders/arcsight_cti.py
@@ -1,7 +1,7 @@
 from app.translator.core.models.platform_details import PlatformDetails
 from app.translator.core.render_cti import RenderCTI
 from app.translator.managers import render_cti_manager
-from app.translator.platforms.arcsight.const import arcsight_query_details, DEFAULT_ARCSIGHT_CTI_MAPPING
+from app.translator.platforms.arcsight.const import DEFAULT_ARCSIGHT_CTI_MAPPING, arcsight_query_details
 
 
 @render_cti_manager.register
diff --git a/uncoder-core/app/translator/platforms/logrhythm_axon/mapping.py b/uncoder-core/app/translator/platforms/logrhythm_axon/mapping.py
@@ -18,9 +18,9 @@ def __str__(self) -> str:
 class LogRhythmAxonMappings(BasePlatformMappings):
     is_strict_mapping = True
 
-    def prepare_mapping(self) -> dict[str, SourceMapping]:
+    def prepare_mapping(self, platform_dir: str) -> dict[str, SourceMapping]:
         source_mappings = {}
-        for mapping_dict in self._loader.load_platform_mappings(self._platform_dir):
+        for mapping_dict in self._loader.load_platform_mappings(platform_dir):
             log_source_signature = self.prepare_log_source_signature(mapping=mapping_dict)
             fields_mapping = self.prepare_fields_mapping(field_mapping=mapping_dict.get("field_mapping", {}))
             source_mappings[DEFAULT_MAPPING_NAME] = SourceMapping(
diff --git a/uncoder-core/app/translator/translator.py b/uncoder-core/app/translator/translator.py
@@ -68,12 +68,19 @@ def parse_meta_info(self, text: str, source: str) -> Union[dict, RawQueryContain
 
     @handle_translation_exceptions
     def __parse_incoming_data(
-        self, text: str, source: str, target: Optional[str] = None
+        self,
+        text: str,
+        source: str,
+        target: Optional[str] = None,
+        source_alt_mapping: Optional[str] = None,
+        target_alt_mapping: Optional[str] = None,
     ) -> tuple[RawQueryContainer, Optional[TokenizedQueryContainer]]:
         parser, raw_query_container = self.parse_raw_query(text=text, source=source)
+        raw_query_container.meta_info.source_alt_mapping = source_alt_mapping
         tokenized_query_container = None
         if not (target and self.__is_one_vendor_translation(raw_query_container.language, target)):
             tokenized_query_container = parser.parse(raw_query_container)
+            tokenized_query_container.meta_info.target_alt_mapping = target_alt_mapping
 
         return raw_query_container, tokenized_query_container
 
@@ -89,8 +96,17 @@ def __render_translation(
             raw_query_container=raw_query_container, tokenized_query_container=tokenized_query_container
         )
 
-    def __translate_one(self, text: str, source: str, target: str) -> (bool, str):
-        status, parsed_data = self.__parse_incoming_data(text=text, source=source, target=target)
+    def __translate_one(
+        self,
+        text: str,
+        source: str,
+        target: str,
+        source_alt_mapping: Optional[str] = None,
+        target_alt_mapping: Optional[str] = None,
+    ) -> (bool, str):
+        status, parsed_data = self.__parse_incoming_data(
+            text=text, source=source, target=target, source_alt_mapping=source_alt_mapping, target_alt_mapping=target_alt_mapping
+        )
         if not status:
             return status, parsed_data
 
@@ -124,10 +140,19 @@ def __translate_all(self, text: str, source: str) -> list[dict]:
 
         return result
 
-    def translate_one(self, text: str, source: str, target: str) -> (bool, str):
+    def translate_one(
+        self,
+        text: str,
+        source: str,
+        target: str,
+        source_alt_mapping: Optional[str] = None,
+        target_alt_mapping: Optional[str] = None,
+    ) -> (bool, str):
         if source == target:
             return True, text
-        return self.__translate_one(text=text, source=source, target=target)
+        return self.__translate_one(
+            text=text, source=source, target=target, source_alt_mapping=source_alt_mapping, target_alt_mapping=target_alt_mapping
+        )
 
     def translate_all(self, text: str, source: str) -> list[dict]:
         return self.__translate_all(text=text, source=source)

Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,3 @@`
`1`		`-from typing import Union`
`2`		`-`
`3`	`1`	`from app.translator.core.const import QUERY_TOKEN_TYPE`
`4`	`2`	`from app.translator.core.custom_types.tokens import LogicalOperatorType, OperatorType`
`5`	`3`	`from app.translator.core.mapping import SourceMapping`