From 880fb300eae9d30c45b5ed88adfff43e72bba12f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rapha=C3=ABl=20Vinot?= <raphael@vinot.info>
Date: Fri, 18 Oct 2019 14:44:54 +0200
Subject: [PATCH] chg: Use default category from template

Fix #477
---
 pymisp/mispevent.py                                | 5 ++++-
 tests/mispevent_testfiles/event_obj_def_param.json | 8 ++++++++
 tests/test_mispevent.py                            | 3 +++
 3 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/pymisp/mispevent.py b/pymisp/mispevent.py
index b99b833..4dbce0a 100644
--- a/pymisp/mispevent.py
+++ b/pymisp/mispevent.py
@@ -1070,6 +1070,9 @@ class MISPObjectAttribute(MISPAttribute):
         self.type = kwargs.pop('type', None)
         if self.type is None:
             self.type = self._definition.get('misp-attribute')
+        if 'category' not in kwargs and 'categories' in self._definition:
+            # Get first category in the list from the object template as default
+            self.category = self._definition['categories'][0]
         self.disable_correlation = kwargs.pop('disable_correlation', None)
         if self.disable_correlation is None:
             # The correlation can be disabled by default in the object definition.
@@ -1331,7 +1334,7 @@ class MISPObject(AbstractMISP):
             logger.warning("The value of the attribute you're trying to add is None or empty string, skipping it. Object relation: {}".format(object_relation))
             return None
         if self._known_template:
-            if self._definition['attributes'].get(object_relation):
+            if object_relation in self._definition['attributes']:
                 attribute = MISPObjectAttribute(self._definition['attributes'][object_relation])
             else:
                 # Woopsie, this object_relation is unknown, no sane defaults for you.
diff --git a/tests/mispevent_testfiles/event_obj_def_param.json b/tests/mispevent_testfiles/event_obj_def_param.json
index ead01d1..5f9907c 100644
--- a/tests/mispevent_testfiles/event_obj_def_param.json
+++ b/tests/mispevent_testfiles/event_obj_def_param.json
@@ -14,6 +14,14 @@
           "to_ids": true,
           "type": "filename",
           "value": "bar"
+        },
+        {
+          "category": "Artifacts dropped",
+          "disable_correlation": false,
+          "object_relation": "pattern-in-file",
+          "to_ids": true,
+          "type": "pattern-in-file",
+          "value": "baz"
         }
       ],
       "description": "File object describing a file with meta-information",
diff --git a/tests/test_mispevent.py b/tests/test_mispevent.py
index 5f31ca4..c0e44b1 100644
--- a/tests/test_mispevent.py
+++ b/tests/test_mispevent.py
@@ -159,6 +159,9 @@ class TestMISPEvent(unittest.TestCase):
         self.mispevent.add_object(name='file', strict=True)
         a = self.mispevent.objects[0].add_attribute('filename', value='bar', Tag=[{'name': 'blah'}])
         del a.uuid
+        a = self.mispevent.objects[0].add_attribute('pattern-in-file', value='baz')
+        self.assertEqual(a.category, 'Artifacts dropped')
+        del a.uuid
         self.mispevent.add_object(name='file', strict=False, default_attributes_parameters=self.mispevent.objects[0].attributes[0])
         a = self.mispevent.objects[1].add_attribute('filename', value='baz')
         del a.uuid