fix cwe_checker_to_ida script (#279)

cwe_checker_to_ida/CweCheckerParser.py

@@ -22,69 +22,37 @@ colors = {'CWE190': YELLOW,
           'CWE787': RED,
           }
 
-class CheckPath(object):
+class Cwe:
 
-    def __init__(self, source, source_addr, destination, destination_addr, path_str):
-        self.source = source
-        self.source_addr = self.__fix_address(source_addr)
-        self.destination = self.__fix_address(destination)
-        self.destination_addr = self.__fix_address(destination_addr)
-        self.path_str = self.__fix_address(path_str)
-        self.color = None
-        self.highlight = False
+    def __init__(self, name, address, description):
+        self.address = '0x'+address
+        self.comment = description
+        self.color = self.__get_color(name)
 
-    @staticmethod
-    def __fix_address(address):
-        return address.replace(':32u', '').replace(':64u', '')
+    def __get_color(self,name):
+        return colors[name]
 
-class CweWarning(object):
+class Parser:
 
-    def __init__(self, name, plugin_version, description, addresses):
-        self.name = name
-        self.plugin_version = plugin_version
-        self.description = self.__fix_address(description)
-        self.color = None
-        self.address = [self.__fix_address(address) for address in addresses]
-        self.highlight = True
-
-    @staticmethod
-    def __fix_address(address):
-        return address.replace(':32u', '').replace(':64u', '')
-
-class Parser(object):
-
-    def __init__(self, result_path):
+    def __init__(self,result_path):
         self._result_path = result_path
-
-    @staticmethod
-    def _parse_cwe_warnings(j):
-        result = []
-
-        if 'warnings' in j:
-            for w in j['warnings']:
-                cwe_warning = CweWarning(w['name'], w['version'], w['description'], w['addresses'])
-                if cwe_warning.name in colors:
-                    cwe_warning.color = colors[cwe_warning.name]
-                else:
-                    cwe_warning.highlight = False
-                result.append(cwe_warning)
-
-        return result
-
-    @staticmethod
-    def _parse_check_path(j):
+    
+    def __parse_cwe(self,j):
         result = []
-
-        if 'check_path' in j:
-            for p in j['check_path']:
-                check_path = CheckPath(p['source'], p['source_addr'], p['destination'], p['destination_addr'], p['path_str'])
-                result.append(check_path)
+        for p in j:
+            addresses = p['addresses']
+            for address in addresses:
+                element = Cwe(
+                        address=address,
+                        name=p['name'],
+                        description=p['description'],
+                )
+                result.append(element)
 
         return result
 
     def parse(self):
         with open(self._result_path) as fhandle:
             j = json.load(fhandle)
-            warnings = self._parse_cwe_warnings(j)
-            check_path = self._parse_check_path(j)
-            return warnings + check_path
+            cwe_out = self.__parse_cwe(j)
+            return cwe_out

cwe_checker_to_ida/CweCheckerParser_test.py

-import unittest
-import json
-import CweCheckerParser
-
-class TestCweCheckerParser(unittest.TestCase):
-
-    def setUp(self):
-        self.parser = CweCheckerParser.Parser('RESULT_PATH')
-
-    def test_parser(self):
-        input_data = json.loads("""{
-        "binary": "test/artificial_samples/build/cwe_190_x86_gcc.out",
-        "time": 1564552342.0,
-        "warnings": [
-        {
-        "name": "CWE190",
-        "version": "0.1",
-        "addresses": [ "0x6BC:32u" ],
-        "symbols": [ "malloc" ],
-        "other": [],
-        "description":
-        "(Integer Overflow or Wraparound) Potential overflow due to multiplication at 0x6BC:32u (malloc)"
-        }]}""")
-        expected_res = 'CWE190'
-
-        res = self.parser._parse_cwe_warnings(input_data)
-
-        self.assertEqual(len(res), 1)
-        self.assertEqual(res[0].name, expected_res)

cwe_checker_to_ida/Generator.py

-from CweCheckerParser import CweWarning
+from CweCheckerParser import Cwe
 
-class IdaGenerator(object):
+class IdaGenerator:
 
     def __init__(self, results):
         self._results = results
@@ -8,17 +8,7 @@ class IdaGenerator(object):
     def generate(self):
         script = "import sark\nimport idaapi\n"
         for res in self._results:
-            if isinstance(res, CweWarning):
-                if res.highlight and res.address:
-                    first_address = res.address[0]
-                    script += "sark.Line(%s).color = %s\n" % (first_address, res.color)
-                    script += "sark.Line(%s).comments.regular = '%s'\n" % (first_address, res.description)
-                    script += "print('[ %s ] %s')\n" % (first_address, res.description)
-                else:
-                    script += "print('[ GENERAL ] %s')\n" % res.description
-            else:
-                script += "print('[CheckPath] %s ( %s ) -> %s via %s')\n" % (res.source,
-                                                                           res.source_addr,
-                                                                           res.destination,
-                                                                           res.path_str)
+            script += "sark.Line(%s).color = %s\n" % (res.address, res.color)
+            script += "sark.Line(%s).comments.regular = '%s'\n" % (res.address, res.comment)
+            script += "print('[ %s ] %s')\n" % (res.address, res.comment)
         return script