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 - name: Test with nosetests
run: | run: |
poetry run pytest --cov=pymispgalaxies tests/tests.py poetry run pytest --cov=pymispgalaxies tests/test*.py
poetry run mypy . poetry run mypy .
- name: Upload coverage to Codecov - name: Upload coverage to Codecov

View File

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

View File

@ -2,7 +2,7 @@
# -*- coding: utf-8 -*- # -*- coding: utf-8 -*-
import unittest import unittest
from pymispgalaxies import Galaxies, Clusters, Cluster, UnableToRevertMachinetag from pymispgalaxies import Galaxies, Clusters, UnableToRevertMachinetag
from glob import glob from glob import glob
import os import os
import json import json
@ -127,13 +127,3 @@ class TestPyMISPGalaxies(unittest.TestCase):
errors[uuid] = entries errors[uuid] = entries
print(json.dumps(errors, indent=2)) print(json.dumps(errors, indent=2))
self.assertFalse(errors) 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')