From effd8084a760089976d31a254beb141e96828312 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rapha=C3=ABl=20Vinot?= <raphael@vinot.info>
Date: Wed, 5 Aug 2015 16:01:57 +0200
Subject: [PATCH] Cleanup of the upload API

---
 examples/upload.py | 13 ++++++++-----
 pymisp/api.py      | 31 +++++++++++++++++++++----------
 2 files changed, 29 insertions(+), 15 deletions(-)

diff --git a/examples/upload.py b/examples/upload.py
index 6954e4f..c4bde86 100755
--- a/examples/upload.py
+++ b/examples/upload.py
@@ -2,13 +2,12 @@
 # -*- coding: utf-8 -*-
 
 from pymisp import PyMISP
-from keys import priv
+from keys import url_priv, key_priv
+# from keys import url_cert, key_cert
 import argparse
 import os
 import glob
 
-url = 'https://misppriv.circl.lu'
-
 
 def init(url, key):
     return PyMISP(url, key, True, 'json')
@@ -16,7 +15,10 @@ def init(url, key):
 
 def upload_files(m, eid, paths, distrib, ids, categ, info, analysis, threat):
     out = m.upload_sample(eid, paths, distrib, ids, categ, info, analysis, threat)
-    print out, out.text
+    if out.status_code == 200:
+        print("Files uploaded sucessfully")
+    else:
+        print("Something failed: {}".format(out.text))
 
 
 if __name__ == '__main__':
@@ -31,7 +33,8 @@ if __name__ == '__main__':
     parser.add_argument("-t", "--threat", type=int, help="The threat level ID of the newly created event, if applicatble. [0-3]")
     args = parser.parse_args()
 
-    misp = init(url, priv)
+    misp = init(url_priv, key_priv)
+    # misp = init(url_cert, key_cert)
 
     files = []
     if os.path.isfile(args.upload):
diff --git a/pymisp/api.py b/pymisp/api.py
index 2e01894..2582dad 100644
--- a/pymisp/api.py
+++ b/pymisp/api.py
@@ -9,6 +9,21 @@ import requests
 import os
 import base64
 
+
+class PyMISPError(Exception):
+    def __init__(self, message):
+        super(PyMISPError, self).__init__(message)
+        self.message = message
+
+
+class NewEventError(PyMISPError):
+    pass
+
+
+class NewAttributeError(PyMISPError):
+    pass
+
+
 class PyMISP(object):
     """
         Python API for MISP
@@ -123,11 +138,11 @@ class PyMISP(object):
     def _create_event(self, distribution, threat_level_id, analysis, info):
         # Setup details of a new event
         if distribution not in [0, 1, 2, 3]:
-            return False
+            raise NewEventError('{} is invalid, the distribution has to be in 0, 1, 2, 3'.format(distribution))
         if threat_level_id not in [0, 1, 2, 3]:
-            return False
+            raise NewEventError('{} is invalid, the threat_level_id has to be in 0, 1, 2, 3'.format(threat_level_id))
         if analysis not in [0, 1, 2]:
-            return False
+            raise NewEventError('{} is invalid, the analysis has to be in 0, 1, 2'.format(analysis))
         return {'distribution': int(distribution), 'info': info,
                 'threat_level_id': int(threat_level_id), 'analysis': analysis}
 
@@ -138,21 +153,17 @@ class PyMISP(object):
             # New event
             postcontent = self._create_event(distribution, threat_level_id,
                                              analysis, info)
-            if postcontent:
-                to_post['request'].update(postcontent)
-            else:
-                # invalid new event
-                return False
+            to_post['request'].update(postcontent)
         else:
             to_post['request'].update({'event_id': int(event_id)})
 
         if to_ids not in [True, False]:
-            return False
+            raise NewAttributeError('{} is invalid, to_ids has to be True or False'.format(analysis))
         to_post['request'].update({'to_ids': to_ids})
 
         if category not in ['Payload delivery', 'Artifacts dropped',
                             'Payload Installation', 'External Analysis']:
-            return False
+            raise NewAttributeError('{} is invalid, category has to be in {}'.format(analysis, (', '.join(['Payload delivery', 'Artifacts dropped', 'Payload Installation', 'External Analysis']))))
         to_post['request'].update({'category': category})
 
         files = []