mirror of https://github.com/MISP/misp-galaxy
96 lines
3.3 KiB
Python
96 lines
3.3 KiB
Python
from dataclasses import dataclass, field, asdict, is_dataclass
|
|
import json
|
|
|
|
@dataclass
|
|
class Meta:
|
|
country: str = None
|
|
country_name: str = None
|
|
refs: list = field(default_factory=list)
|
|
synonyms: list = field(default_factory=list)
|
|
|
|
# def custom_asdict(obj):
|
|
# if is_dataclass(obj):
|
|
# result = {}
|
|
# for field_name, field_def in obj.__dataclass_fields__.items():
|
|
# value = getattr(obj, field_name)
|
|
# if field_name == 'meta':
|
|
# meta_value = custom_asdict(value)
|
|
# meta_value = {k: v for k, v in meta_value.items() if not (k in ['refs', 'synonyms'] and (not v or all(e is None for e in v)))}
|
|
# value = meta_value
|
|
# elif isinstance(value, (list, tuple)) and all(is_dataclass(i) for i in value):
|
|
# value = [custom_asdict(i) for i in value]
|
|
# elif isinstance(value, list) and all(e is None for e in value):
|
|
# continue
|
|
# result[field_name] = value
|
|
# return result
|
|
# else:
|
|
# return obj
|
|
|
|
def custom_asdict(obj):
|
|
if is_dataclass(obj):
|
|
result = {}
|
|
for field_name, field_def in obj.__dataclass_fields__.items():
|
|
value = getattr(obj, field_name)
|
|
if field_name == 'meta':
|
|
meta_value = custom_asdict(value)
|
|
# Filter out 'refs', 'synonyms', 'country', and 'country_name' if they are None or if 'refs' and 'synonyms' are empty
|
|
meta_value = {k: v for k, v in meta_value.items() if v is not None and not (k in ['refs', 'synonyms'] and (not v or all(e is None for e in v)))}
|
|
value = meta_value
|
|
elif isinstance(value, (list, tuple)) and all(is_dataclass(i) for i in value):
|
|
value = [custom_asdict(i) for i in value]
|
|
elif isinstance(value, list) and all(e is None for e in value):
|
|
continue
|
|
# Skip the field if the value is None (specifically for 'country' and 'country_name')
|
|
if value is None and field_name in ['country', 'country_name']:
|
|
continue
|
|
result[field_name] = value
|
|
return result
|
|
else:
|
|
return obj
|
|
|
|
@dataclass
|
|
class IntelAgency:
|
|
description: str = ""
|
|
meta: Meta = field(default_factory=Meta)
|
|
related: list = field(default_factory=list)
|
|
uuid: str = None
|
|
value: str = None
|
|
|
|
def __post_init__(self):
|
|
if not self.value:
|
|
raise ValueError("IntelAgency 'value' cannot be empty.")
|
|
if not self.uuid:
|
|
raise ValueError("IntelAgency 'uuid' cannot be empty.")
|
|
|
|
@dataclass
|
|
class Galaxy:
|
|
description: str
|
|
icon: str
|
|
name: str
|
|
namespace: str
|
|
type: str
|
|
uuid: str
|
|
version: int
|
|
|
|
def save_to_file(self, path: str):
|
|
with open(path, "w") as file:
|
|
file.write(json.dumps(asdict(self), indent=4))
|
|
|
|
@dataclass
|
|
class Cluster:
|
|
authors: str
|
|
category: str
|
|
description: str
|
|
name: str
|
|
source: str
|
|
type: str
|
|
uuid: str
|
|
version: int
|
|
values: list = field(default_factory=list)
|
|
|
|
def add_value(self, value: IntelAgency):
|
|
self.values.append(value)
|
|
|
|
def save_to_file(self, path: str):
|
|
with open(path, "w") as file:
|
|
file.write(json.dumps(custom_asdict(self), indent=4, ensure_ascii=False)) |