#39 Allow specifying custom fields when creating a case

samples/test-case-create.py

@@ -9,19 +9,32 @@
 import json
 import time
 from thehive4py.api import TheHiveApi
-from thehive4py.models import Case, CaseTask
+from thehive4py.models import Case, CaseTask, CustomFieldHelper
 
 api = TheHiveApi('http://127.0.0.1:9000', 'username', 'password', {'http': '', 'https': ''})
 
-
 # Prepare the sample case
 tasks = [
     CaseTask(title='Tracking'),
     CaseTask(title='Communication'),
     CaseTask(title='Investigation', status='Waiting', flag=True)
 ]
-# tasks = []
-case = Case(title='From TheHive4Py', tlp=3, flag=True, tags=['TheHive4Py', 'sample'], description='N/A', tasks=tasks)
+
+# Prepare the custom fields
+customFields = CustomFieldHelper()\
+    .add_boolean('booleanField', True)\
+    .add_string('businessImpact', 'HIGH')\
+    .add_date('occurDate', int(time.time())*1000)\
+    .add_number('cvss', 9)\
+    .build()
+
+case = Case(title='From TheHive4Py',
+            tlp=3,
+            flag=True,
+            tags=['TheHive4Py', 'sample'],
+            description='N/A',
+            tasks=tasks,
+            customFields=customFields)
 
 # Create the case
 print('Create Case')

thehive4py/models.py

@@ -34,6 +34,36 @@ def attr(self, attributes, name, default, error=None):
             return attributes.get(name, default)
 
 
+class CustomFieldHelper(object):
+    def __init__(self):
+        self.fields = {}
+
+    def __add_field(self, type, name, value):
+        custom_field = dict()
+        custom_field['order'] = len(self.fields)
+        custom_field[type] = value
+        self.fields[name] = custom_field
+
+    def add_date(self, name, value):
+        self.__add_field('date', name, value)
+        return self
+
+    def add_string(self, name, value):
+        self.__add_field('string', name, value)
+        return self
+
+    def add_boolean(self, name, value):
+        self.__add_field('boolean', name, value)
+        return self
+
+    def add_number(self, name, value):
+        self.__add_field('number', name, value)
+        return self
+
+    def build(self):
+        return self.fields
+
+
 class Case(JSONSerializable):
 
     def __init__(self, **attributes):
@@ -46,6 +76,7 @@ def __init__(self, **attributes):
             'tags': [],
             'startDate': int(time.time()) * 1000,
             'metrics': {},
+            'customFields': {},
             'tasks': [],
             'template': None
         }
@@ -65,6 +96,7 @@ def __init__(self, **attributes):
         self.tags = attributes.get('tags', defaults['tags'])
         self.startDate = attributes.get('startDate', defaults['startDate'])
         self.metrics = attributes.get('metrics', defaults['metrics'])
+        self.customFields = attributes.get('customFields', defaults['customFields'])
         self.template = attributes.get('template', defaults['template'])
 
         tasks = attributes.get('tasks', defaults['tasks'])
@@ -177,7 +209,8 @@ def __init__(self, **attributes):
         self.flag = attributes.get('flag', False)
         self.tlp = attributes.get('tlp', 2)
         self.tags = attributes.get('tags', [])
-        self.metricNames = attributes.get('metricNames', [])
+        self.metrics = attributes.get('metrics', {})
+        self.customFields = attributes.get('customFields', {})
 
         tasks = attributes.get('tasks', [])
         self.tasks = []