chg: [Cluster] add(cv: ClusterValue)

pull/25/head
Christophe Vandeplas 2024-06-18 10:36:15 +02:00
parent 8c2f69dbfd
commit b5de7b54d4
No known key found for this signature in database
GPG Key ID: BDC48619FFDC5A5B
4 changed files with 42 additions and 20 deletions

View File

@ -33,7 +33,7 @@ jobs:
- name: Test with nosetests
run: |
poetry run pytest --cov=pymispgalaxies tests/tests.py
poetry run pytest --cov=pymispgalaxies tests/test*.py
poetry run mypy .
- name: Upload coverage to Codecov

View File

@ -394,14 +394,12 @@ class Cluster(Mapping): # type: ignore
self.category = self.cluster['category']
self.cluster_values = {}
self.duplicates = []
for value in self.cluster['values']:
new_cluster_value = ClusterValue(value)
if self.get(new_cluster_value.value):
if skip_duplicates:
self.duplicates.append((self.name, new_cluster_value.value))
else:
raise PyMISPGalaxiesError("Duplicate value ({}) in cluster: {}".format(new_cluster_value.value, self.name))
self.cluster_values[new_cluster_value.value] = new_cluster_value
try:
for value in self.cluster['values']:
new_cluster_value = ClusterValue(value)
self.add(new_cluster_value, skip_duplicates)
except KeyError:
pass
@overload
def search(self, query: str, return_tags: Literal[False] = False) -> List[ClusterValue]:
@ -465,6 +463,17 @@ class Cluster(Mapping): # type: ignore
return value
raise KeyError('No value with external_id: {}'.format(external_id))
def add(self, cv: ClusterValue, skip_duplicates: bool) -> None:
"""
Adds a cluster value to the cluster.
"""
if self.get(cv.value):
if skip_duplicates:
self.duplicates.append((self.name, cv.value))
else:
raise PyMISPGalaxiesError("Duplicate value ({}) in cluster: {}".format(cv.value, self.name))
self.cluster_values[cv.value] = cv
def __str__(self) -> str:
"""
Returns a string representation of the cluster.

View File

@ -2,7 +2,7 @@
# -*- coding: utf-8 -*-
import unittest
from pymispgalaxies import Galaxies, Clusters, Cluster, UnableToRevertMachinetag
from pymispgalaxies import Galaxies, Clusters, UnableToRevertMachinetag
from glob import glob
import os
import json
@ -127,13 +127,3 @@ class TestPyMISPGalaxies(unittest.TestCase):
errors[uuid] = entries
print(json.dumps(errors, indent=2))
self.assertFalse(errors)
def test_get_by_external_id(self):
cluster = Cluster(cluster='mitre-attack-pattern')
self.assertIsNotNone(cluster)
cluster_by_external_id = cluster.get_by_external_id('T1525')
cluster_by_value = cluster.get('Implant Internal Image - T1525')
self.assertEqual(cluster_by_external_id, cluster_by_value)
with self.assertRaises(KeyError):
cluster.get_by_external_id('XXXXXX')

23
tests/tests_api.py Normal file
View File

@ -0,0 +1,23 @@
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import unittest
from pymispgalaxies import Galaxies, Clusters, Cluster
class TestPyMISPGalaxiesApi(unittest.TestCase):
def setUp(self):
self.galaxies = Galaxies()
self.clusters = Clusters(skip_duplicates=False)
self.maxDiff = None
def test_get_by_external_id(self):
cluster = Cluster(cluster='mitre-attack-pattern')
self.assertIsNotNone(cluster)
cluster_by_external_id = cluster.get_by_external_id('T1525')
cluster_by_value = cluster.get('Implant Internal Image - T1525')
self.assertEqual(cluster_by_external_id, cluster_by_value)
with self.assertRaises(KeyError):
cluster.get_by_external_id('XXXXXX')