Merge pull request #79 from ANSSI-FR/dev

Bugfixes and IDA Pro typedefs

examples/reschedule_analysis.py

@@ -2,8 +2,8 @@
 
 from time import sleep
 
-from app import app
-from app.controllers.api import APIControl
+from poli import app
+from poli.controllers.api import APIControl
 
 
 def launch():

examples/schedule_samples_analysis.py

@@ -2,9 +2,9 @@
 import sys
 from time import sleep
 
-from app.controllers.api import APIControl
+from poli.controllers.api import APIControl
 api = APIControl()
 
 for i in sys.argv[1:]:
-    api.analysiscontrol.schedule_sample_analysis(int(i))
+    api.analysiscontrol.schedule_sample_analysis(int(i), force=True)
     sleep(200)

poli/controllers/analysis.py

@@ -144,8 +144,6 @@ def reschedule_all_analysis(self, force=False):
         for sample in Sample.query.all():
             if force or sample.analysis_status == AnalysisStatus.TOSTART:
                 self.schedule_sample_analysis(sample.id, force)
-            elif sample.analysis_status == AnalysisStatus.RUNNING:
-                self.schedule_sample_analysis(sample.id, force)
 
 
 class Analysis(object):

poli/controllers/idaactions.py

@@ -15,6 +15,7 @@
 from poli.models.idaactions import IDANameAction, IDACommentAction
 from poli.models.idaactions import IDAStruct, IDAStructSchema
 from poli.models.idaactions import IDAStructMember
+from poli.models.idaactions import IDATypeAction
 from poli.models.sample import Sample
 
 
@@ -207,3 +208,25 @@ def change_struct_member_size(struct_id, mid, new_size):
         db.session.commit()
 
         return True
+
+    @staticmethod
+    def add_typedef(address, typedef):
+        """
+            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()
+        return mtype.id
+
+    @classmethod
+    def get_typedefs(cls, sid, addr=None, timestamp=None):
+        """
+            Return filtered IDA Pro type definitions
+        """
+        data = cls.filter_actions('idatypes', sid, addr, timestamp)
+        schema = IDAActionSchema(many=True)
+        return schema.dump(data).data

poli/controllers/sample.py

@@ -318,8 +318,9 @@ def search_hash(cls, needle):
         c = Sample.query.filter(Sample.md5.like(needleq)).all()
         results = list(set(a + b + c))
         function_results = None
-        if re.match("[0-9a-f]{8}", needle):
-            function_results = cls.search_machoc_single_hash(needle)
+	# XXX fix this
+        #if re.match("[0-9a-f]{8}", needle):
+            #function_results = cls.search_machoc_single_hash(needle)
         return results, function_results
 
     @classmethod
@@ -432,23 +433,28 @@ def match_by_importhash(cls, sample):
                     cls.add_sample_match(sample, sample_2, "iat_hash")
                     # add the corresponding match to the other sample
                     cls.add_sample_match(sample_2, sample, "iat_hash")
-                continue
         return True
 
+    @staticmethod
+    def query_machoc_matches(sample,  sample_2):
+        query = SampleMatch.query.filter(SampleMatch.sid_1.in_([sample.id, sample_2.id]), SampleMatch.sid_2.in_(
+                [sample.id, sample_2.id]), SampleMatch.match_type == "machoc80")
+        if query.count() != 0:
+            return True
+        return False
+
     @classmethod
     def match_by_machoc80(cls, sample):
         """
             Match samples by machoc hash.
         """
-        if len(sample.functions) == 0:
+        if sample.functions.count() == 0:
             return True
         for sample_2 in Sample.query.all():
-            if len(sample_2.functions) == 0:
-                continue
-            if SampleMatch.query.filter(SampleMatch.sid_1.in_([sample.id, sample_2.id]), SampleMatch.sid_2.in_(
-                    [sample.id, sample_2.id]), SampleMatch.match_type == "machoc80").count() != 0:
+            if cls.query_machoc_matches(sample, sample_2):
                 continue
-            if cls.machoc_diff_samples(sample, sample_2) >= 0.8:
+            elif cls.machoc_diff_samples(sample, sample_2) >= 0.8:
+                app.logger.debug("Add machoc match %d %d", sample.id, sample_2.id)
                 cls.add_sample_match(sample, sample_2, "machoc80")
                 cls.add_sample_match(sample_2, sample, "machoc80")
         return True
@@ -458,11 +464,11 @@ def machoc_diff_with_all_samples(cls, sample, level=0.8):
         """
             Diff a sample with all other samples. Class method.
         """
-        if len(sample.functions) == 0:
+        if sample.functions.count() == 0:
             return []
         hits = []
         for sample_2 in Sample.query.all():
-            if len(sample_2.functions) == 0 or sample_2 == sample:
+            if sample_2.functions.count() == 0 or sample_2.id == sample.id:
                 continue
             hit_rate = cls.machoc_diff_samples(sample, sample_2)
             if hit_rate >= level:
@@ -474,8 +480,6 @@ def machoc_diff_samples(cls, sample1, sample2):
         """
             Diff two samples using machoc.
         """
-        if sample1 == sample2:
-            return 0
         sample1_hashes = []
         sample2_hashes = []
         for f in sample1.functions:
@@ -484,7 +488,8 @@ def machoc_diff_samples(cls, sample1, sample2):
         for f in sample2.functions:
             if f.machoc_hash is not None and f.machoc_hash != -1:
                 sample2_hashes.append(f.machoc_hash)
-        return cls.machoc_diff_hashes(sample1_hashes, sample2_hashes)
+        rate = cls.machoc_diff_hashes(sample1_hashes, sample2_hashes)
+        return rate
 
     @staticmethod
     def machoc_diff_hashes(sample1_hashes, sample2_hashes):
@@ -496,7 +501,8 @@ def machoc_diff_hashes(sample1_hashes, sample2_hashes):
         maxlen = max(len(sample1_hashes), len(sample2_hashes))
         c1, c2 = map(Counter, (sample1_hashes, sample2_hashes))
         ch = set(sample1_hashes).intersection(set(sample2_hashes))
-        return float(sum(map(lambda h: max(c1[h], c2[h]), ch))) / maxlen
+        rate = float(sum(map(lambda h: max(c1[h], c2[h]), ch))) / maxlen
+        return rate
 
     def machoc_get_similar_functions(self, sample_dst, sample_src):
         """
@@ -686,7 +692,15 @@ def add_multiple_strings(self, sample, strings):
         return True
 
     @staticmethod
-    def add_function(sample, address, machoc_hash,
+    def query_function_info(sample, address):
+        obj = FunctionInfo.query.filter_by(sample_id=sample.id, address=address)
+        if obj.count() != 0:
+            return obj.first()
+        else:
+            return None
+
+    @classmethod
+    def add_function(cls, sample, address, machoc_hash,
                      name="", overwrite=False, do_commit=True):
         """
             Add a function. Updates if exists.
@@ -696,9 +710,8 @@ def add_function(sample, address, machoc_hash,
         if name == "":
             name = "sub_" + hex(address)[2:]
         functions_exists = False
-        obj = FunctionInfo.query.filter_by(sample=sample, address=address)
-        if obj.count() != 0:
-            function_info = obj.first()
+        function_info = cls.query_function_info(sample, address)
+        if function_info is not None:
             functions_exists = True
             if not overwrite:
                 return True

poli/controllers/tasks/task_analyzeitrb.py

@@ -52,7 +52,7 @@ def get_addr_data(line):
         addr, data = None, None
         if len(items) == 3:
             addr = int(items[1], 16)
-            data = items[2][:-1]
+            data = items[2].replace('\n','')
         return addr, data
 
     def parse_machoc_signatures(self):

poli/models/family.py

@@ -88,16 +88,16 @@ def fromstring(cls, val):
 # Yara signatures relationship (auto-classification).
 familytoyara = db.Table('familytoyara',
                         db.Column('yara_id', db.Integer,
-                                  db.ForeignKey('yararule.id')),
+                                  db.ForeignKey('yararule.id'), index=True),
                         db.Column('family_id', db.Integer,
-                                  db.ForeignKey('family.id'))
+                                  db.ForeignKey('family.id'), index=True)
                         )
 # Samples relationship.
 familytosample = db.Table('familytosample',
                           db.Column('sample_id', db.Integer,
-                                    db.ForeignKey('sample.id')),
+                                    db.ForeignKey('sample.id'), index=True),
                           db.Column('family_id', db.Integer,
-                                    db.ForeignKey('family.id'))
+                                    db.ForeignKey('family.id'), index=True)
                           )
 
 
@@ -117,7 +117,7 @@ class Family(db.Model):
                             secondary=familytoyara,
                             backref=db.backref('families', lazy='dynamic'))
     # 1-N relationships
-    parent_id = db.Column(db.Integer, db.ForeignKey('family.id'))
+    parent_id = db.Column(db.Integer, db.ForeignKey('family.id'), index=True)
     subfamilies = db.relationship(
         'Family', backref=db.backref(
             'parents', remote_side=[id]))

poli/models/idaactions.py

@@ -55,7 +55,9 @@ class IDACommentAction(IDAAction):
 
 
 class IDANameAction(IDAAction):
-    # Represent names in IDA
+    """
+        This represents global names in IDA.
+    """
     __tablename__ = 'idanames'
     id = db.Column(db.Integer(),
                    db.ForeignKey('idaactions.id'),
@@ -64,6 +66,19 @@ class IDANameAction(IDAAction):
         'polymorphic_identity': 'idanames'}
 
 
+class IDATypeAction(IDAAction):
+    """
+    This represents the types as applied by
+    the shortcut 'Y' in IDA Pro
+    """
+    __tablename__ = 'idatypes'
+    id = db.Column(db.Integer(),
+                   db.ForeignKey('idaactions.id'),
+                   primary_key=True)
+    __mapper_args__ = {
+        'polymorphic_identity': 'idatypes'}
+
+
 class IDAApplyStructs(IDAAction):
     # This is the action of applying a structure to an address
     __tablename__ = 'idaapplystructs'
@@ -116,6 +131,7 @@ class Meta:
             "type",
         )
 
+
 class IDAStructMemberSchema(ma.ModelSchema):
     class Meta:
         fields = (

poli/models/sample.py

@@ -277,7 +277,7 @@ class Sample(db.Model):
             'sample', remote_side=[id]))
     functions = db.relationship(
         "FunctionInfo", backref=db.backref(
-            'sample', remote_side=[id]))
+            'sample', remote_side=[id]), lazy="dynamic")
     filenames = db.relationship(
         "FileName", backref=db.backref(
             'sample', remote_side=[id]))

poli/views/apiview.py

@@ -292,7 +292,8 @@ def api_get_sample_id_from_hash(shash):
         abort(400, "Invalid hash length")
     if sample is not None:
         return jsonify({'sample_id': sample.id})
-    return jsonify({'sample_id': None})
+    else:
+        abort(404)
 
 
 @apiview.route('/samples/<int:sid>/download/')
@@ -534,6 +535,29 @@ def api_post_sample_names(sid):
     return jsonify({'result': result})
 
 
+@apiview.route('/samples/<int:sid>/types/', methods=['POST'])
+def api_post_sample_types(sid):
+    """
+        Manage the creation of type definitions at specific addresses
+    """
+    data = request.json
+    addr = data['address']
+    typedef = data['typedef']
+
+    action_id = api.idacontrol.add_typedef(addr, typedef)
+    result = api.samplecontrol.add_idaaction(sid, action_id)
+    return jsonify(dict(result=result))
+
+@apiview.route('/samples/<int:sid>/types/', methods=['GET'])
+def api_get_sample_types(sid):
+    """
+        Get the IDA types stored in DB
+    """
+    current_timestamp, addr = get_filter_arguments(request)
+    data = api.idacontrol.get_typedefs(sid, addr, current_timestamp)
+    return jsonify({'typedefs': data})
+
+
 @apiview.route('/samples/<int:sid>/structs/', methods=['POST'])
 def api_create_struct(sid):
     """

poli/views/webui.py

@@ -65,7 +65,7 @@ def before_request():
         Affects global variables for each request
     """
     g.user = current_user
-    g.samples = Sample.query.all()
+    g.samples = Sample.query.order_by(Sample.id).limit(15).all()
 
 
 """

poliapi/mainapi.py

@@ -159,6 +159,19 @@ def set_abstract(self, sid, abstract):
         json_data = dict(abstract=abstract)
         return self.post(endpoint, json=json_data)["result"]
 
+    def get_abstract(self, sid):
+        endpoint = self.prepare_endpoint(root='samples')
+        endpoint += str(sid) + '/abstract/'
+
+        return self.get(endpoint)["abstract"]
+
+
+    def get_sid_from_MD5(self, md5):
+        endpoint = self.prepare_endpoint(root='samples')
+        endpoint += md5 + '/'
+
+        return self.get(endpoint)
+
 
 class FamilyModule(MainModule):
     """