chg: allow to pass a list of galaxies as parameter

pull/3/head
Raphaël Vinot 2018-02-23 13:46:48 +01:00
parent c51e562a4a
commit d76c6082ab
2 changed files with 22 additions and 14 deletions

View File

@ -22,6 +22,7 @@ class EncodeGalaxies(JSONEncoder):
return obj.to_dict() return obj.to_dict()
return JSONEncoder.default(self, obj) return JSONEncoder.default(self, obj)
class EncodeClusters(JSONEncoder): class EncodeClusters(JSONEncoder):
def default(self, obj): def default(self, obj):
if isinstance(obj, (Cluster, ClusterValue, ClusterValueMeta)): if isinstance(obj, (Cluster, ClusterValue, ClusterValueMeta)):
@ -60,13 +61,17 @@ class Galaxy():
class Galaxies(collections.Mapping): class Galaxies(collections.Mapping):
def __init__(self): def __init__(self, galaxies=[]):
self.root_dir_galaxies = os.path.join(os.path.abspath(os.path.dirname(sys.modules['pymispgalaxies'].__file__)), if not galaxies:
'data', 'misp-galaxy', 'galaxies') galaxies = []
self.root_dir_galaxies = os.path.join(os.path.abspath(os.path.dirname(sys.modules['pymispgalaxies'].__file__)),
'data', 'misp-galaxy', 'galaxies')
for galaxy_file in glob(os.path.join(self.root_dir_galaxies, '*.json')):
with open(galaxy_file, 'r') as f:
galaxies.append(json.load(f))
self.galaxies = {} self.galaxies = {}
for galaxy_file in glob(os.path.join(self.root_dir_galaxies, '*.json')): for galaxy in galaxies:
with open(galaxy_file, 'r') as f:
galaxy = json.load(f)
self.galaxies[galaxy['name']] = Galaxy(galaxy) self.galaxies[galaxy['name']] = Galaxy(galaxy)
def validate_with_schema(self): def validate_with_schema(self):
@ -239,13 +244,16 @@ class Cluster(collections.Mapping):
class Clusters(collections.Mapping): class Clusters(collections.Mapping):
def __init__(self): def __init__(self, clusters=[]):
self.root_dir_clusters = os.path.join(os.path.abspath(os.path.dirname(sys.modules['pymispgalaxies'].__file__)), if not clusters:
'data', 'misp-galaxy', 'clusters') clusters = []
self.root_dir_clusters = os.path.join(os.path.abspath(os.path.dirname(sys.modules['pymispgalaxies'].__file__)),
'data', 'misp-galaxy', 'clusters')
for cluster_file in glob(os.path.join(self.root_dir_clusters, '*.json')):
with open(cluster_file, 'r') as f:
clusters.append(json.load(f))
self.clusters = {} self.clusters = {}
for cluster_file in glob(os.path.join(self.root_dir_clusters, '*.json')): for cluster in clusters:
with open(cluster_file, 'r') as f:
cluster = json.load(f)
self.clusters[cluster['type']] = Cluster(cluster) self.clusters[cluster['type']] = Cluster(cluster)
def validate_with_schema(self): def validate_with_schema(self):
@ -267,7 +275,7 @@ class Clusters(collections.Mapping):
cluster = self.get(cluster_type) cluster = self.get(cluster_type)
value = cluster[cluster_value] value = cluster[cluster_value]
return cluster, value return cluster, value
except: except Exception:
raise UnableToRevertMachinetag('The machinetag {} could not be found.'.format(machinetag)) raise UnableToRevertMachinetag('The machinetag {} could not be found.'.format(machinetag))
def search(self, query): def search(self, query):

@ -1 +1 @@
Subproject commit 9cfba28d3612621943b282ca434b26845f011c55 Subproject commit 0eb6ab483280c943bfb3a6c19ce5ba1320f25c83