From 5062c6162099efa296a81690d17849ed398b3353 Mon Sep 17 00:00:00 2001 From: niclas Date: Fri, 23 Feb 2024 14:54:25 +0100 Subject: [PATCH] Add [tidal] relation enrichment with mitre --- tools/tidal-api/main.py | 16 ++- tools/tidal-api/models/cluster.py | 168 ++++++++++++++++++++++++------ 2 files changed, 149 insertions(+), 35 deletions(-) diff --git a/tools/tidal-api/main.py b/tools/tidal-api/main.py index 9438567..559a4f5 100644 --- a/tools/tidal-api/main.py +++ b/tools/tidal-api/main.py @@ -17,7 +17,7 @@ GALAXY_PATH = "../../galaxies" CLUSTER_PATH = "../../clusters" -def create_galaxy(endpoint: str, version: int): +def create_galaxy(endpoint: str, version: int, extended_relations: bool = False): api = TidalAPI() data = api.get_data(endpoint) with open(f"{CONFIG}/{endpoint}.json", "r") as file: @@ -28,10 +28,10 @@ def create_galaxy(endpoint: str, version: int): match endpoint: case "groups": - cluster = GroupCluster(**config["cluster"], uuid=galaxy.uuid) + cluster = GroupCluster(**config["cluster"], uuid=galaxy.uuid, enrichment=extended_relations) cluster.add_values(data) case "software": - cluster = SoftwareCluster(**config["cluster"], uuid=galaxy.uuid) + cluster = SoftwareCluster(**config["cluster"], uuid=galaxy.uuid, enrichment=extended_relations) cluster.add_values(data) case "campaigns": cluster = CampaignsCluster(**config["cluster"], uuid=galaxy.uuid) @@ -56,9 +56,9 @@ def create_galaxy(endpoint: str, version: int): def main(args, galaxies): if args.all: for galaxy in galaxies: - create_galaxy(galaxy, args.version) + create_galaxy(galaxy, args.version, args.extended_relations) else: - create_galaxy(args.type, args.version) + create_galaxy(args.type, args.version, args.extended_relations) if __name__ == "__main__": @@ -72,6 +72,7 @@ if __name__ == "__main__": description="Create galaxy and cluster json files from Tidal API" ) parser.add_argument( + "-a", "--all", action="store_true", help="Create all galaxies and clusters", @@ -88,6 +89,11 @@ if __name__ == "__main__": required=True, help="The version of the galaxy", ) + parser.add_argument( + "--extended-relations", + action="store_true", + help="Create extended relations in the cluster", + ) parser.set_defaults(func=main) args = parser.parse_args() diff --git a/tools/tidal-api/models/cluster.py b/tools/tidal-api/models/cluster.py index fca9171..f85839d 100644 --- a/tools/tidal-api/models/cluster.py +++ b/tools/tidal-api/models/cluster.py @@ -1,4 +1,5 @@ from dataclasses import dataclass, field, asdict +from typing import Type import json @@ -19,6 +20,13 @@ class GroupsMeta(Meta): owner: str = None +@dataclass +class AssociatedGroupsMeta(Meta): + id: str = None + owner_id: str = None + owner: str = None + + @dataclass class SoftwareMeta(Meta): source: str = None @@ -29,6 +37,13 @@ class SoftwareMeta(Meta): owner: str = None +@dataclass +class AssociatedSoftwareMeta(Meta): + id: str = None + owner_id: str = None + owner: str = None + + @dataclass class TechniqueMeta(Meta): source: str = None @@ -108,9 +123,10 @@ class Cluster: self.type = type self.uuid = uuid self.values = [] + self.CLUSTER_PATH = "../../clusters" - def add_values(self): - print("This method should be implemented in the child class") + def add_values(self, data: dict, meta_class: Type[Meta]): + pass def save_to_file(self, path): with open(path, "w") as file: @@ -131,6 +147,24 @@ class Cluster: "values": self.values, } + def _get_relation_from_mitre_id( + self, mitre_id: str, cluster: str, meta_key: str, array: bool = False + ): + with open(f"{self.CLUSTER_PATH}/{cluster}.json", "r") as file: + mitre = json.load(file) + for entry in mitre["values"]: + try: + if array: + for id in entry["meta"][meta_key]: + if id == mitre_id: + return entry["uuid"] + else: + if entry["meta"][meta_key] == mitre_id: + return entry["uuid"] + except KeyError: + continue + return None + class GroupCluster(Cluster): def __init__( @@ -142,8 +176,10 @@ class GroupCluster(Cluster): source: str, type: str, uuid: str, + enrichment: bool = False, ): super().__init__(authors, category, description, name, source, type, uuid) + self.enrichment = enrichment def add_values(self, data): for entry in data["data"]: @@ -166,13 +202,39 @@ class GroupCluster(Cluster): owner=entry.get("owner_name"), ) related = [] - for relation in entry.get("associated_groups"): + if self.enrichment: + related_cluster = self._get_relation_from_mitre_id( + entry.get("group_attack_id"), "threat-actor", "synonyms", True + ) + if related_cluster: + related.append( + { + "dest-uuid": related_cluster, + "type": "similar", + } + ) + + for associated_group in entry.get("associated_groups"): + meta = AssociatedGroupsMeta( + id=associated_group.get("id"), + owner_id=associated_group.get("owner_id"), + owner=associated_group.get("owner_name"), + ) + value = ClusterValue( + description=associated_group.get("description"), + meta=meta, + related=[], + uuid=associated_group.get("associated_group_id"), + value=associated_group.get("name"), + ) + self.values.append(value.return_value()) related.append( { - "dest-uuid": relation.get("id"), - "type": "related-to", + "dest-uuid": associated_group.get("associated_group_id"), + "type": "similar", } ) + value = ClusterValue( description=entry.get("description"), meta=meta, @@ -193,8 +255,10 @@ class SoftwareCluster(Cluster): source: str, type: str, uuid: str, + enrichment: bool = False, ): super().__init__(authors, category, description, name, source, type, uuid) + self.enrichment = enrichment def add_values(self, data): for entry in data["data"]: @@ -214,13 +278,50 @@ class SoftwareCluster(Cluster): "type": "used-by", } ) - for relation in entry.get("associated_software"): + if self.enrichment: + related_cluster = self._get_relation_from_mitre_id( + entry.get("software_attack_id"), "mitre-tool", "external_id" + ) + if related_cluster: + related.append( + { + "dest-uuid": related_cluster, + "type": "similar", + } + ) + + related_cluster = self._get_relation_from_mitre_id( + entry.get("software_attack_id"), "mitre-malware", "external_id" + ) + if related_cluster: + related.append( + { + "dest-uuid": related_cluster, + "type": "similar", + } + ) + + for associated_software in entry.get("associated_software"): + meta = AssociatedSoftwareMeta( + id=associated_software.get("id"), + owner_id=associated_software.get("owner_id"), + owner=associated_software.get("owner_name"), + ) + value = ClusterValue( + description=associated_software.get("description"), + meta=meta, + related=[], + uuid=associated_software.get("associated_software_id"), + value=associated_software.get("name"), + ) + self.values.append(value.return_value()) related.append( { - "dest-uuid": relation.get("id"), - "type": "related-to", + "dest-uuid": associated_software.get("associated_software_id"), + "type": "similar", } ) + value = ClusterValue( description=entry.get("description"), meta=meta, @@ -260,6 +361,35 @@ class TechniqueCluster(Cluster): "type": "uses", } ) + + for sub_technique in entry.get("sub_technique"): + meta = SubTechniqueMeta( + source=sub_technique.get("source"), + technique_attack_id=sub_technique.get("technique_attack_id"), + ) + sub_related = [] + for relation in sub_technique.get("tactic"): + sub_related.append( + { + "dest-uuid": relation.get("tactic_id"), + "type": "uses", + } + ) + sub_value = ClusterValue( + description=sub_technique.get("description"), + meta=meta, + related=sub_related, + uuid=sub_technique.get("id"), + value=sub_technique.get("name"), + ) + self.values.append(sub_value.return_value()) + related.append( + { + "dest-uuid": sub_technique.get("id"), + "type": "similar", + } + ) + value = ClusterValue( description=entry.get("description"), meta=meta, @@ -269,28 +399,6 @@ class TechniqueCluster(Cluster): ) self.values.append(value.return_value()) - for sub_technique in entry.get("sub_technique"): - meta = SubTechniqueMeta( - source=sub_technique.get("source"), - technique_attack_id=sub_technique.get("technique_attack_id"), - ) - related = [] - for relation in sub_technique.get("tactic"): - related.append( - { - "dest-uuid": relation.get("tactic_id"), - "type": "uses", - } - ) - value = ClusterValue( - description=sub_technique.get("description"), - meta=meta, - related=related, - uuid=sub_technique.get("id"), - value=sub_technique.get("name"), - ) - self.values.append(value.return_value()) - class TacticCluster(Cluster): def __init__(