Updated AssociatedObject for CybOX 2.0

Author: ikiril01

cybox/core/associated_object.py

@@ -1,9 +1,11 @@
 # Copyright (c) 2013, The MITRE Corporation. All rights reserved.
 # See LICENSE.txt for complete terms.
 
+import cybox
 import cybox.utils as utils
-import cybox.bindings.cybox_core_1_0 as core_binding
+import cybox.bindings.cybox_core as core_binding
 from cybox.core.object import Object
+from cybox.common import VocabString
 
 class AssociatedObject(Object):
 
@@ -18,7 +20,7 @@ def __init__(self, type_=None, association_type_=None, defined_object=None):
 
     def to_obj(self):
         obj = super(AssociatedObject, self).to_obj(core_binding.AssociatedObjectType())
-        obj.set_association_type(self.association_type_)
+        obj.set_Association_Type(self.association_type_.to_obj())
         return obj
 
     def to_dict(self):
@@ -31,29 +33,13 @@ def from_obj(object_obj):
         if not object_obj:
             return None
         obj = Object.from_obj(object_obj, AssociatedObject())
-        obj.association_type_ = object_obj.get_association_type()
+        obj.association_type_ = VocabString.from_obj(object_obj.get_association_type())
         return obj
 
     @staticmethod
     def from_dict(object_dict):
         if not object_dict:
             return None
         obj = Object.from_dict(object_dict, AssociatedObject())
-        obj.association_type_ = object_dict.get('association_type', None)
-        return obj
-
-
-    #@classmethod
-    #def object_from_dict(cls, associated_object_dict):
-    #    """Create the Associated Object Python object representation from an input dictionary"""
-    #    associated_obj = Object.object_from_dict(associated_object_dict,core_binding.AssociatedObjectType())
-    #    for key, value in associated_object_dict.items():
-    #        if key == 'association_type' and utils.test_value(value): associated_obj.set_association_type(value)
-    #    return associated_obj
-
-    #@classmethod
-    #def dict_from_object(cls, associated_obj):
-    #    """Parse and return a dictionary for an Associated Object"""
-    #    associated_object_dict = Object.dict_from_object(associated_obj)
-    #    if associated_obj.get_association_type() is not None: associated_object_dict['association_type'] = associated_obj.get_association_type()
-    #    return associated_object_dict
+        obj.association_type_ = VocabString.from_dict(object_dict.get('association_type', None))
+        return obj