Merge branch 'release/1.2.0'

samples/sample.txt

@@ -0,0 +1 @@
+TheHive4Py

samples/test-alert-create.py

@@ -0,0 +1,58 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import requests
+import sys
+import json
+import time
+import uuid
+from thehive4py.api import TheHiveApi
+from thehive4py.models import Alert, AlertArtifact
+
+api = TheHiveApi('http://127.0.0.1:9000', 'username', 'password', {'http': '', 'https': ''})
+
+artifacts = [
+    AlertArtifact(dataType='ip', data='8.8.8.8'),
+    AlertArtifact(dataType='domain', data='google.com'),
+    AlertArtifact(dataType='file', data='pic.png'),
+    AlertArtifact(dataType='file', data='sample.txt')
+]
+
+
+# Prepare the sample Alert
+sourceRef = str(uuid.uuid4())[0:6]
+alert = Alert(title='New Alert',
+              tlp=3,
+              tags=['TheHive4Py', 'sample'],
+              description='N/A',
+              type='external',
+              source='instance1',
+              sourceRef=sourceRef,
+              artifacts=artifacts)
+
+# Create the Alert
+print('Create Alert')
+print('-----------------------------')
+id = None
+response = api.create_alert(alert)
+if response.status_code == 201:
+    print(json.dumps(response.json(), indent=4, sort_keys=True))
+    print('')
+    id = response.json()['id']
+else:
+    print('ko: {}/{}'.format(response.status_code, response.text))
+    sys.exit(0)
+
+
+# Get all the details of the created alert
+print('Get created alert {}'.format(id))
+print('-----------------------------')
+response = api.get_alert(id)
+if response.status_code == requests.codes.ok:
+    print(json.dumps(response.json(), indent=4, sort_keys=True))
+    print('')
+else:
+    print('ko: {}/{}'.format(response.status_code, response.text))

setup.py

@@ -11,7 +11,7 @@
 
 setup(
     name='thehive4py',
-    version='1.1.1',
+    version='1.2.0',
     description='Python API client for TheHive.',
     long_description=read_md('README.md'),
     author='TheHive-Project',

thehive4py/api.py

@@ -223,6 +223,54 @@ def get_case_template(self, name):
         except requests.exceptions.RequestException as e:
             sys.exit("Error: {}".format(e))
 
+    def get_case_tasks(self, caseId):
+        req = self.url + "/api/case/task/_search?range=all"
+        data = {
+            'query': {
+                '_parent': {
+                    '_type': 'case',
+                    '_query': {
+                        '_id': caseId
+                    }
+                }
+            }}
+        try:
+            response = requests.post(req, json=data, proxies=self.proxies, auth=self.auth)
+            if response.status_code == 200 and len(response.json()) > 0:
+                return response.json()
+            else:
+                sys.exit("Error: {}".format("Unable to find tasks"))
+        except requests.exceptions.RequestException as e:
+            sys.exit("Error: {}".format(e))
+
+    def create_alert(self, alert):
+
+        """
+        :param alert: TheHive alert
+        :type alert: Alert defined in models.py
+        :return: TheHive alert
+        :rtype: json
+        """
+
+        req = self.url + "/api/alert"
+        data = alert.jsonify()
+        try:
+            return requests.post(req, headers={'Content-Type': 'application/json'}, data=data, proxies=self.proxies, auth=self.auth)
+        except requests.exceptions.RequestException as e:
+            sys.exit("Error: {}".format(e))
+
+    def get_alert(self, alert_id):
+        """
+            :param alert_id: Alert identifier
+            :return: TheHive Alert
+            :rtype: json
+        """
+        req = self.url + "/api/alert/{}".format(alert_id)
+
+        try:
+            return requests.get(req, proxies=self.proxies, auth=self.auth)
+        except requests.exceptions.RequestException as e:
+            sys.exit("Error: {}".format(e))
 
 # - addObservable(file)
 # - addObservable(data)

thehive4py/models.py

@@ -5,6 +5,9 @@
 import time
 import os
 import magic
+import base64
+
+from future.utils import raise_with_traceback
 
 
 class CustomJsonEncoder(json.JSONEncoder):
@@ -19,6 +22,14 @@ class JSONSerializable(object):
     def jsonify(self):
         return json.dumps(self, sort_keys=True, indent=4, cls=CustomJsonEncoder)
 
+    def attr(self, attributes, name, default, error=None):
+        is_required = error is not None
+
+        if is_required and name not in attributes:
+            raise_with_traceback(ValueError(error))
+        else:
+            return attributes.get(name, default)
+
 
 class Case(JSONSerializable):
 
@@ -120,6 +131,7 @@ def __init__(self, **attributes):
             else:
                 self.tasks.append(CaseTask(json=task))
 
+
 class CaseObservable(JSONSerializable):
     def __init__(self, **attributes):
         if attributes.get('json', False):
@@ -132,7 +144,56 @@ def __init__(self, **attributes):
 
         data = attributes.get('data', [])
         if self.dataType == 'file':
-
-            self.data = [{'attachment': ( os.path.basename(data[0]), open(data[0], 'rb'), magic.Magic(mime=True).from_file(data[0]))}]
+            self.data = [{'attachment': (os.path.basename(data[0]), open(data[0], 'rb'), magic.Magic(mime=True).from_file(data[0]))}]
         else:
             self.data = data
+
+
+class Alert(JSONSerializable):
+    def __init__(self, **attributes):
+        if attributes.get('json', False):
+            attributes = attributes['json']
+
+        self.tlp = attributes.get('tlp', 2)
+        self.severity = attributes.get('severity', 3)
+        self.date = attributes.get('date', int(time.time()) * 1000)
+        self.tags = attributes.get('tags', [])
+        self.caseTemplate = attributes.get('caseTemplate', None)
+
+        self.title = self.attr(attributes, 'title', None, 'Missing alert title')
+        self.type = self.attr(attributes, 'type', None, 'Missing alert type')
+        self.source = self.attr(attributes, 'source', None, 'Missing alert source')
+        self.sourceRef = self.attr(attributes, 'sourceRef', None, 'Missing alert reference')
+        self.description = self.attr(attributes, 'description', None, 'Missing alert description')
+
+        artifacts = attributes.get('artifacts', [])
+        self.artifacts = []
+        for artifact in artifacts:
+            if type(artifact) == AlertArtifact:
+                self.artifacts.append(artifact)
+            else:
+                self.artifacts.append(AlertArtifact(json=artifact))
+
+
+class AlertArtifact(JSONSerializable):
+    def __init__(self, **attributes):
+        if attributes.get('json', False):
+            attributes = attributes['json']
+
+        self.dataType = attributes.get('dataType', None)
+        self.message = attributes.get('message', None)
+        self.tlp = attributes.get('tlp', 2)
+        self.tags = attributes.get('tags', [])
+
+        if self.dataType == 'file':
+            self.data = self._prepare_file_data(attributes.get('data', None))
+        else:
+            self.data = attributes.get('data', None)
+
+    def _prepare_file_data(self, file_path):
+        with open(file_path, "rb") as file_artifact:
+            filename = os.path.basename(file_path)
+            mime = magic.Magic(mime=True).from_file(file_path)
+            encoded_string = base64.b64encode(file_artifact.read())
+
+        return filename + ';' + mime + ';' + encoded_string;