Merge pull request #163 from ANSSI-FR/fix-skelenox-auth

Fix skelenox auth

poli/controllers/idaactions.py

@@ -1,8 +1,7 @@
 """
     This file is part of Polichombr.
 
-    (c) 2017 ANSSI-FR
-
+    (c) 2018 ANSSI-FR
 
     Description:
         Managers for all the actions associated with IDAPro models
@@ -26,6 +25,9 @@ class IDAActionsController(object):
     """
     @staticmethod
     def get_all(sid=None, timestamp=None):
+        """
+            Return all the actions for a sample from a timestamp
+        """
         if sid is None:
             return False
         query = IDAAction.query
@@ -37,16 +39,25 @@ def get_all(sid=None, timestamp=None):
         return schema.dump(query.all()).data
 
     @staticmethod
-    def add_comment(address, data):
+    def setup_generic_action(action, address, data, user_id=None):
+        """
+            Setup and commit the common parts of an IDA action
+        """
+        action.address = address
+        action.data = data
+        action.timestamp = datetime.datetime.now()
+        action.user_id = user_id
+        db.session.add(action)
+        db.session.commit()
+        return action
+
+    @classmethod
+    def add_comment(cls, address, data, user_id=None):
         """
             Creates a new comment action
         """
         comment = IDACommentAction()
-        comment.address = address
-        comment.data = data
-        comment.timestamp = datetime.datetime.now()
-        db.session.add(comment)
-        db.session.commit()
+        comment = cls.setup_generic_action(comment, address, data, user_id)
         return comment.id
 
     @staticmethod
@@ -78,17 +89,13 @@ def get_comments(cls, sid=None, addr=None, timestamp=None):
         data = schema.dump(data).data
         return data
 
-    @staticmethod
-    def add_name(address=None, data=None):
+    @classmethod
+    def add_name(cls, address=None, data=None, user_id=None):
         """
             Creates a new name action
         """
         name = IDANameAction()
-        name.address = address
-        name.data = data
-        name.timestamp = datetime.datetime.now()
-        db.session.add(name)
-        db.session.commit()
+        name = cls.setup_generic_action(name, address, data, user_id)
         return name.id
 
     @classmethod
@@ -102,22 +109,28 @@ def get_names(cls, sid, addr=None, timestamp=None):
         schema = IDAActionSchema(many=True)
         return schema.dump(data).data
 
-    @staticmethod
-    def create_struct(name=None):
+    @classmethod
+    def create_struct(cls, name=None, user_id=None):
+        """
+            Create a new struct.
+        """
         if name is None:
             app.logger.error("Cannot create anonymous struct")
             return False
         mstruct = IDAStruct()
+        mstruct = cls.setup_generic_action(mstruct, 0, name, user_id)
+
         mstruct.name = name
-        mstruct.data = name
-        mstruct.timestamp = datetime.datetime.now()
         mstruct.size = 0
         db.session.add(mstruct)
         db.session.commit()
         return mstruct.id
 
     @staticmethod
     def get_structs(sid, timestamp=None):
+        """
+            List all structs from a given timestamp
+        """
         query = IDAStruct.query
         query = query.filter(IDAStruct.samples.any(Sample.id == sid))
 
@@ -166,7 +179,7 @@ def rename_struct(struct_id, name):
     @staticmethod
     def delete_struct(struct_id):
         """
-
+            Delete a struct, but not it's members
         """
         mstruct = IDAStruct.query.get(struct_id)
         app.logger.debug("Deleting struct %s", mstruct.name)
@@ -176,6 +189,9 @@ def delete_struct(struct_id):
 
     @staticmethod
     def create_struct_member(name=None, size=None, offset=None):
+        """
+            Create a struct member, but don't affect it to the struct yet
+        """
         member = IDAStructMember()
         if member is None:
             return False
@@ -241,17 +257,13 @@ def change_struct_member_size(cls, struct_id, mid, new_size):
 
         return True
 
-    @staticmethod
-    def add_typedef(address, typedef):
+    @classmethod
+    def add_typedef(cls, address, typedef, user_id):
         """
             Creates a new type definition
         """
         mtype = IDATypeAction()
-        mtype.address = address
-        mtype.data = typedef
-        mtype.timestamp = datetime.datetime.now()
-        db.session.add(mtype)
-        db.session.commit()
+        mtype = cls.setup_generic_action(mtype, address, typedef, user_id)
         return mtype.id
 
     @classmethod

poli/controllers/sample.py

@@ -55,7 +55,7 @@ def create_sample_from_file(self, file_data, orig_filename="", user=None,
         # If not, we will just update some information
         if Sample.query.filter_by(sha256=sha_256).count() != 0:
             sample = Sample.query.filter_by(sha256=sha_256).first()
-            if sample.storage_file is not None and sample.storage_file != "" and os.path.exists(
+            if sample.storage_file is not None and os.path.exists(
                     sample.storage_file):
                 return sample
 
@@ -637,12 +637,12 @@ def add_metadata(sample, metadata_type, metadata_value):
             metadata_value = hex(metadata_value)
         else:
             try:
-                metadata_value = str(metadata_value)
+                metadata_value = str(metadata_value).replace("\x00", "")
             except Exception as e:
                 app.logger.exception(e)
                 return False
-        for s_metadata in sample.s_metadata:
-            if s_metadata.type_id == metadata_type and s_metadata.value == metadata_value:
+        for meta in sample.s_metadata:
+            if meta.type_id == metadata_type and meta.value == metadata_value:
                 return True
         s_metadata = SampleMetadata()
         s_metadata.value = metadata_value

poli/templates/sample.html

@@ -231,7 +231,7 @@ <h3 class="panel-title">IDA PRO commands <a class="btn btn-xs btn-default" style
                                             <td>{{"0x%x"|format((act.address|int))}}</td>
                                             <td>{{act.type}}</td>
                                             <td>{{act.data}}</td>
-                                            <td>Analyst</td>
+											<td>{{act.user_id}}</td>
 					    <td><button type="button" class="warning" aria-label="Close">Delete</span></button> </td>
                                         </tr>
                                     {% endfor %}

poli/views/api_idaactions.py

@@ -10,7 +10,7 @@
 import datetime
 
 from flask import jsonify, request, abort, current_app
-from flask_security import login_required
+from flask_security import login_required, current_user
 
 from poli import api
 from poli.views.apiview import apiview
@@ -105,12 +105,13 @@ def api_post_sample_comments(sid):
         abort(400, "Missing comment or address arguments")
     address = data['address']
     comment = data['comment']
+    user_id = current_user.id
     current_app.logger.debug(
         "Getting a new comment for sample %d : %s@0x%x",
         sid,
         comment,
         address)
-    action_id = api.idacontrol.add_comment(address, comment)
+    action_id = api.idacontrol.add_comment(address, comment, user_id)
     result = api.samplecontrol.add_idaaction(sid, action_id)
     return jsonify({'result': result})
 
@@ -141,12 +142,13 @@ def api_post_sample_names(sid):
     data = request.json
     addr = data['address']
     name = data['name']
+    user_id = current_user.id
     current_app.logger.debug(
         "Getting a new name for sample %d : %s@0x%x",
         sid,
         name,
         addr)
-    action_id = api.idacontrol.add_name(addr, name)
+    action_id = api.idacontrol.add_name(addr, name, user_id)
     result = api.samplecontrol.add_idaaction(sid, action_id)
     if result is True:
         api.samplecontrol.rename_func_from_action(sid, addr, name)
@@ -164,8 +166,9 @@ def api_post_sample_types(sid):
     data = request.json
     addr = data['address']
     typedef = data['typedef']
+    user_id = current_user.id
 
-    action_id = api.idacontrol.add_typedef(addr, typedef)
+    action_id = api.idacontrol.add_typedef(addr, typedef, user_id)
     result = api.samplecontrol.add_idaaction(sid, action_id)
     return jsonify(dict(result=result))
 
@@ -194,7 +197,8 @@ def api_create_struct(sid):
     result = False
     name = data['name']
     current_app.logger.debug("Creating structure %s" % name)
-    mstruct = api.idacontrol.create_struct(name=name)
+    user_id = current_user.id
+    mstruct = api.idacontrol.create_struct(name=name, user_id=user_id)
     if mstruct is not False:
         result = api.samplecontrol.add_idaaction(sid, mstruct)
     return jsonify({'result': result, 'structs': [{'id': mstruct}]})

poli/views/apiview.py

@@ -14,7 +14,7 @@
 from poli.models.user import User
 
 from flask import jsonify, request, abort, make_response
-from flask import Blueprint
+from flask import Blueprint, current_app
 
 apiview = Blueprint('apiview', __name__,
                     url_prefix=app.config['API_PATH'])
@@ -106,7 +106,8 @@ def generate_token():
     key = data['api_key']
     user = User.query.filter_by(api_key=key).first()
     if not user:
-        abort(400, "Invalid user")
+        current_app.logger.error("Invalid user trying to login")
+        abort(401, "Invalid user")
     return jsonify({'token': user.get_auth_token()})
 
 

run.py

@@ -1,7 +1,18 @@
 #!/usr/bin/env python
+"""
+    This file is part of Polichombr.
+
+    (c) 2018 ANSSI-FR
+
+
+    Description:
+        Wrapper for running the application
+"""
+
 from poli import app
 
 if __name__ == "__main__":
+
     app.run(app.config['SERVER_ADDR'],
             port=app.config['SERVER_PORT'],
             debug=app.config['DEBUG'])

skelenox.py

@@ -131,6 +131,7 @@ def __init__(self, skel_config=None):
         self.poli_port = skel_config.poli_port
 
         self.h_conn = None
+        self.auth_token = None
         self.is_online = False
         self.sample_id = None
 
@@ -156,6 +157,7 @@ def __do_init(self):
 
         self.h_conn = httplib.HTTPConnection(self.poli_server, self.poli_port)
         self.h_conn.connect()
+        self.login()
         self.is_online = True
         self.init_sample_id()
 
@@ -169,6 +171,24 @@ def close_connection(self):
         self.is_online = False
         self.sample_id = None
 
+    def login(self):
+        data = json.dumps({'api_key': self.api_key})
+        headers = {"Accept-encoding": "gzip, deflate",
+                   "Content-type": "application/json",
+                   "Accept": "*/*;q=0.8",
+                   "Accept-Language": "en-US,en;q=0.5"}
+
+        self.h_conn.request("POST",
+                            "/api/1.0/auth_token/",
+                            data,
+                            headers)
+        res = self.h_conn.getresponse()
+        if res.status != 200:
+            idc.warning("Error, cannot login to Polichombr!")
+            raise IOError
+        token = json.loads(res.read())["token"]
+        self.auth_token = token
+
     def poli_request(self, endpoint, data, method="POST"):
         """
             @arg : endpoint The API target endpoint
@@ -183,7 +203,7 @@ def poli_request(self, endpoint, data, method="POST"):
                    "Accept": "*/*;q=0.8",
                    "Accept-Language": "en-US,en;q=0.5",
                    "Connection": "Keep-Alive",
-                   "X-API-Key": self.api_key}
+                   "X-Api-Key": self.auth_token}
         json_data = json.dumps(data)
         try:
             self.h_conn.request(method, endpoint, json_data, headers)
@@ -194,7 +214,11 @@ def poli_request(self, endpoint, data, method="POST"):
             self.h_conn.request(method, endpoint, json_data, headers)
         res = self.h_conn.getresponse()
 
-        if res.status != 200:
+        if res.status == 401:
+            g_logger.error("Token is invalid, trying to login again")
+            self.login()
+            return None
+        elif res.status != 200:
             g_logger.error("The %s request didn't go as expected", method)
             g_logger.debug("Status code was %d and content was %s",
                            res.status, res.read())
@@ -283,7 +307,7 @@ def send_sample(self, filedata):
         """
         endpoint = "/api/1.0/samples/"
         headers = {"Accept-encoding": "gzip, deflate",
-                   "X-API-Key": self.api_key}
+                   "X-Api-Key": self.auth_token}
 
         method = "POST"
         boundary = "70f6e331562f4b8f98e5f9590e0ffb8e"

tests/tests_api.py

@@ -673,7 +673,7 @@ def test_get_all(self):
         self.assertIn(actions["idaactions"][3]["data"], "TESTCOMMENT2")
 
         self.assertIn(actions["idaactions"][4]["type"], "idastructs")
-        self.assertEqual(actions["idaactions"][4]["address"], None)
+        self.assertEqual(actions["idaactions"][4]["address"], 0)
         self.assertIn(actions["idaactions"][4]["data"], "ThisIsAStruct")
 
     def test_push_comments(self):