Skip to content

Commit

Permalink
Code cleanup.
Browse files Browse the repository at this point in the history
* Replaced several list fields with EntityList and TypedList
  implementations.
* Added new EntityList types.
  • Loading branch information
Bryan Worrell committed Mar 1, 2015
1 parent ee3515f commit 8fd207f
Show file tree
Hide file tree
Showing 11 changed files with 142 additions and 128 deletions.
2 changes: 1 addition & 1 deletion stix/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@

__version__ = "1.1.1.3"

from .base import Entity, EntityList
from .base import Entity, EntityList, TypedList

# Make sure common gets imported before anything else.
from . import common
36 changes: 10 additions & 26 deletions stix/campaign/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,24 +36,8 @@ class Attribution(GenericRelationshipList):
_inner_name = "threat_actors"


class AttributionList(stix.EntityList):
# NOT AN ACTUAL STIX CLASS. DO NOT CALL `.to_obj()`
# ON THIS DIRECTLY! THIS IS BEING USED FOR CASTING
# PURPOSES ONLY.
_namespace = "http://stix.mitre.org/Campaign-1"
_binding = None
_binding_class = None
_binding_var = None
class AttributionList(stix.TypedList):
_contained_type = Attribution
_inner_name = None

def _fix_value(self, value):
try:
new_value = self._contained_type(None, value)
except:
raise ValueError("Can't put '%s' (%s) into a %s" %
(value, type(value), self.__class__))
return new_value


class RelatedIncidents(GenericRelationshipList):
Expand Down Expand Up @@ -250,11 +234,14 @@ def attribution(self):

@attribution.setter
def attribution(self, value):
if isinstance(value, AttributionList):
self._attribution = value
return

self._attribution = AttributionList()

if not value:
return
elif isinstance(value, AttributionList):
self._attribution = value
elif hasattr(value, '__getitem__'):
self._attribution = AttributionList(*value)
else:
Expand Down Expand Up @@ -290,7 +277,7 @@ def to_obj(self, return_obj=None, ns_info=None):
if self.related_indicators:
return_obj.Related_Indicators = self.related_indicators.to_obj(ns_info=ns_info)
if self.attribution:
return_obj.Attribution = [x.to_obj(ns_info=ns_info) for x in self.attribution]
return_obj.Attribution = self.attribution.to_obj(ns_info=ns_info)
if self.associated_campaigns:
return_obj.Associated_Campaigns = self.associated_campaigns.to_obj(ns_info=ns_info)
if self.confidence:
Expand Down Expand Up @@ -332,8 +319,7 @@ def from_obj(cls, obj, return_obj=None):
RelatedIncidents.from_obj(obj.Related_Incidents)
return_obj.related_indicators = \
RelatedIndicators.from_obj(obj.Related_Indicators)
return_obj.attribution = \
[Attribution.from_obj(x) for x in obj.Attribution]
return_obj.attribution = AttributionList.from_obj(obj.Attribution)
return_obj.associated_campaigns = \
AssociatedCampaigns.from_obj(obj.Associated_Campaigns)
return_obj.confidence = Confidence.from_obj(obj.Confidence)
Expand Down Expand Up @@ -376,7 +362,7 @@ def to_dict(self):
if self.related_indicators:
d['related_indicators'] = self.related_indicators.to_dict()
if self.attribution:
d['attribution'] = [x.to_dict() for x in self.attribution]
d['attribution'] = self.attribution.to_list()
if self.associated_campaigns:
d['associated_campaigns'] = self.associated_campaigns.to_dict()
if self.confidence:
Expand Down Expand Up @@ -419,9 +405,7 @@ def from_dict(cls, dict_repr, return_obj=None):
RelatedIncidents.from_dict(dict_repr.get('related_incidents'))
return_obj.related_indicators = \
RelatedIndicators.from_dict(dict_repr.get('related_indicators'))
return_obj.attribution = \
[Attribution.from_dict(x) for x in
dict_repr.get('attribution', [])]
return_obj.attribution = AttributionList.from_list(dict_repr.get('attribution'))
return_obj.associated_campaigns = \
AssociatedCampaigns.from_dict(dict_repr.get('associated_campaigns'))
return_obj.confidence = \
Expand Down

0 comments on commit 8fd207f

Please sign in to comment.