misp-galaxy/tools/tidal-api/main.py

132 lines
3.6 KiB
Python
Raw Normal View History

2024-02-21 16:24:48 +01:00
from api.api import TidalAPI
from models.galaxy import Galaxy
2024-02-23 11:25:07 +01:00
from models.cluster import (
GroupCluster,
SoftwareCluster,
CampaignsCluster,
TechniqueCluster,
TacticCluster,
ReferencesCluster,
)
2024-02-21 16:24:48 +01:00
import argparse
2024-02-23 11:14:00 +01:00
import json
import os
2024-02-21 16:24:48 +01:00
2024-02-23 11:14:00 +01:00
CONFIG = "./config"
GALAXY_PATH = "../../galaxies"
CLUSTER_PATH = "../../clusters"
2024-02-21 16:24:48 +01:00
2024-02-23 11:25:07 +01:00
2024-03-05 14:27:25 +01:00
def create_galaxy(
endpoint: str,
version: int,
extended_relations: bool = False,
create_subs: bool = False,
):
2024-02-21 16:24:48 +01:00
api = TidalAPI()
2024-02-23 11:14:00 +01:00
data = api.get_data(endpoint)
with open(f"{CONFIG}/{endpoint}.json", "r") as file:
config = json.load(file)
galaxy = Galaxy(**config["galaxy"], version=version)
galaxy.save_to_file(f"{GALAXY_PATH}/tidal-{endpoint}.json")
2024-02-23 11:25:07 +01:00
2024-02-23 11:14:00 +01:00
match endpoint:
case "groups":
2024-03-05 14:27:25 +01:00
cluster = GroupCluster(
**config["cluster"],
uuid=galaxy.uuid,
enrichment=extended_relations,
subs=create_subs,
2024-03-05 15:33:10 +01:00
version=version,
2024-03-05 14:27:25 +01:00
)
2024-02-23 11:14:00 +01:00
cluster.add_values(data)
case "software":
2024-03-05 14:27:25 +01:00
cluster = SoftwareCluster(
**config["cluster"],
uuid=galaxy.uuid,
2024-03-05 15:33:10 +01:00
version=version,
2024-03-05 14:27:25 +01:00
enrichment=extended_relations,
subs=create_subs,
)
2024-02-23 11:14:00 +01:00
cluster.add_values(data)
case "campaigns":
2024-03-05 15:33:10 +01:00
cluster = CampaignsCluster(**config["cluster"], uuid=galaxy.uuid, version=version)
2024-02-23 11:14:00 +01:00
cluster.add_values(data)
case "technique":
2024-03-05 14:27:25 +01:00
cluster = TechniqueCluster(
2024-03-05 15:33:10 +01:00
**config["cluster"], uuid=galaxy.uuid, subs=create_subs, version=version
2024-03-05 14:27:25 +01:00
)
2024-02-23 11:14:00 +01:00
cluster.add_values(data)
case "tactic":
2024-03-05 15:33:10 +01:00
cluster = TacticCluster(**config["cluster"], uuid=galaxy.uuid, version=version)
2024-02-23 11:14:00 +01:00
cluster.add_values(data)
case "references":
2024-03-05 15:33:10 +01:00
cluster = ReferencesCluster(**config["cluster"], uuid=galaxy.uuid, version=version)
2024-02-23 11:14:00 +01:00
cluster.add_values(data)
case _:
print("Error: Invalid endpoint")
return
cluster.save_to_file(f"{CLUSTER_PATH}/tidal-{endpoint}.json")
print(f"Galaxy tidal-{endpoint} created")
2024-02-23 11:25:07 +01:00
2024-02-23 11:14:00 +01:00
def main(args, galaxies):
2024-02-21 16:24:48 +01:00
if args.all:
2024-02-23 11:14:00 +01:00
for galaxy in galaxies:
2024-03-05 14:27:25 +01:00
create_galaxy(
galaxy, args.version, args.extended_relations, args.create_subs
)
2024-02-21 16:24:48 +01:00
else:
2024-03-05 14:27:25 +01:00
create_galaxy(
args.type, args.version, args.extended_relations, args.create_subs
)
2024-02-23 11:25:07 +01:00
2024-02-22 10:52:23 +01:00
2024-02-21 16:24:48 +01:00
if __name__ == "__main__":
2024-02-23 11:14:00 +01:00
galaxies = []
for f in os.listdir(CONFIG):
if f.endswith(".json"):
galaxies.append(f.split(".")[0])
2024-02-22 10:52:23 +01:00
parser = argparse.ArgumentParser(
2024-02-23 11:14:00 +01:00
description="Create galaxy and cluster json files from Tidal API"
2024-02-22 10:52:23 +01:00
)
2024-02-23 11:14:00 +01:00
parser.add_argument(
"-a",
2024-02-23 11:14:00 +01:00
"--all",
action="store_true",
help="Create all galaxies and clusters",
2024-02-22 10:52:23 +01:00
)
2024-02-23 11:14:00 +01:00
parser.add_argument(
2024-02-22 10:52:23 +01:00
"--type",
2024-02-23 11:14:00 +01:00
choices=galaxies,
help="The type of the file to create",
2024-02-22 10:52:23 +01:00
)
2024-02-23 11:14:00 +01:00
parser.add_argument(
"-v",
"--version",
type=int,
required=True,
help="The version of the galaxy",
)
parser.add_argument(
"--extended-relations",
action="store_true",
2024-03-05 14:18:15 +01:00
help="Create extended relations for the clusters",
)
parser.add_argument(
"--create-subs",
action="store_true",
help="Create subclusters from the API",
)
2024-02-23 11:14:00 +01:00
parser.set_defaults(func=main)
2024-02-21 16:24:48 +01:00
args = parser.parse_args()
2024-02-22 10:52:23 +01:00
if hasattr(args, "func"):
2024-02-23 11:14:00 +01:00
args.func(args, galaxies=galaxies)
2024-02-21 16:24:48 +01:00
else:
2024-02-23 11:25:07 +01:00
parser.print_help()