ENH: added __iadd__ to MaltegoTransformRequestMessage

malleum-inc · Aug 8, 2016 · 5222bc5 · 5222bc5
1 parent 831c4e1
commit 5222bc5
Showing 1 changed file with 13 additions and 1 deletion.
diff --git a/src/canari/maltego/message.py b/src/canari/maltego/message.py
@@ -38,7 +38,8 @@
     'TimeSpanEntityField',
     'RegexEntityField',
     'ColorEntityField',
-    'Entity'
+    'Entity',
+    'Limits'
 ]
 
 
@@ -630,6 +631,17 @@ class MaltegoTransformRequestMessage(MaltegoElement):
     _parameters = fields_.Dict(Field, tagname='TransformFields', key='name', required=False)
     limits = fields_.Model(Limits, required=False)
 
+    def __iadd__(self, other):
+        if isinstance(other, Entity):
+            self.__entities.append(other.__entity__)
+        elif isinstance(other, _Entity):
+            self.__entities.append(other)
+        elif isinstance(other, Field):
+            self._parameters[other.name] = other
+        elif isinstance(other, Limits):
+            self.limits = other
+        return self
+
     @property
     def entity(self):
         """Returns the first Entity object in the transform request.