fix: [sharinggroups] Fixes wrong model for SharingGroupOrg

pull/810/head
Christophe Vandeplas 2021-11-30 10:38:41 +01:00
parent 167438ba13
commit afedd8d90b
1 changed files with 41 additions and 15 deletions

View File

@ -125,35 +125,61 @@ class MISPOrganisation(AbstractMISP):
return f'<{self.__class__.__name__}(NotInitialized)' return f'<{self.__class__.__name__}(NotInitialized)'
class MISPSharingGroupOrg(AbstractMISP):
_fields_for_feed: set = {'extend', 'Organisation'}
def __init__(self):
super().__init__()
self.extend: bool
self.Organisation: MISPOrganisation
def from_dict(self, **kwargs):
if 'SharingGroupOrg' in kwargs:
kwargs = kwargs['SharingGroupOrg']
if 'Organisation' in kwargs:
self.Organisation = MISPOrganisation()
self.Organisation.from_dict(**kwargs.pop('Organisation'))
super().from_dict(**kwargs)
def __repr__(self) -> str:
if hasattr(self, 'Organisation') and hasattr(self, 'extend'):
return f'<{self.__class__.__name__}(Org={self.Organisation.name}, extend={self.extend})'
return f'<{self.__class__.__name__}(NotInitialized)'
def _to_feed(self) -> Dict:
to_return = super()._to_feed()
to_return['Organisation'] = self.Organisation._to_feed()
return to_return
class MISPSharingGroup(AbstractMISP): class MISPSharingGroup(AbstractMISP):
_fields_for_feed: set = {'uuid', 'name', 'roaming', 'created', 'organisation_uuid', 'Organisation', 'SharingGroupOrg', 'SharingGroupServer'} _fields_for_feed: set = {'uuid', 'name', 'roaming', 'created', 'organisation_uuid', 'Organisation', 'SharingGroupOrg', 'SharingGroupServer'}
def __init__(self): def __init__(self):
super().__init__() super().__init__()
self.name: str self.name: str
self.SharingGroupOrg: List[MISPOrganisation] = [] self.SharingGroupOrg: List[MISPSharingGroupOrg] = []
@property @property
def orgs(self) -> List[MISPOrganisation]: def sgorgs(self) -> List[MISPSharingGroupOrg]:
return self.SharingGroupOrg return self.SharingGroupOrg
@orgs.setter @sgorgs.setter
def orgs(self, orgs: List[MISPOrganisation]): def sgorgs(self, sgorgs: List[MISPSharingGroupOrg]):
"""Set a list of prepared MISPSighting.""" if all(isinstance(x, MISPSharingGroupOrg) for x in sgorgs):
if all(isinstance(x, MISPSighting) for x in orgs): self.SharingGroupOrg = sgorgs
self.SharingGroupOrg = orgs
else: else:
raise PyMISPError('All the attributes have to be of type MISPOrganisation.') raise PyMISPError('All the attributes have to be of type MISPSharingGroupOrg.')
def add_org(self, org): def add_sgorg(self, sgorg):
misp_org = MISPOrganisation() misp_sgorg = MISPSharingGroupOrg()
misp_org.from_dict(**org) misp_sgorg.from_dict(**sgorg)
self.SharingGroupOrg.append(misp_org) self.SharingGroupOrg.append(misp_sgorg)
return misp_org return misp_sgorg
def from_dict(self, **kwargs): def from_dict(self, **kwargs):
if 'SharingGroupOrg' in kwargs: if 'SharingGroupOrg' in kwargs:
[self.add_org(org) for org in kwargs.pop('SharingGroupOrg')] [self.add_sgorg(sgorg) for sgorg in kwargs.pop('SharingGroupOrg')]
if 'SharingGroup' in kwargs: if 'SharingGroup' in kwargs:
kwargs = kwargs['SharingGroup'] kwargs = kwargs['SharingGroup']
super().from_dict(**kwargs) super().from_dict(**kwargs)
@ -165,7 +191,7 @@ class MISPSharingGroup(AbstractMISP):
def _to_feed(self) -> Dict: def _to_feed(self) -> Dict:
to_return = super()._to_feed() to_return = super()._to_feed()
to_return['SharingGroupOrg'] = [org._to_feed() for org in self.SharingGroupOrg] to_return['SharingGroupOrg'] = [sgorg._to_feed() for sgorg in self.SharingGroupOrg]
to_return['Organisation'].pop('id', None) to_return['Organisation'].pop('id', None)
for server in to_return['SharingGroupServer']: for server in to_return['SharingGroupServer']:
server.pop('id', None) server.pop('id', None)