Dongtai engine/filter

dongtai_common/engine/vul_engine.py

@@ -196,6 +196,7 @@ def loop(self, index, size, current_link):
                     current_link.append(self.copy_method(sub_method, source=True))
                     self.vul_source_signature = f"{sub_method.get('className')}.{sub_method.get('methodName')}"
                     self.vul_stack.append(current_link[::-1])
+                    self.taint_value = sub_method['targetValues']
                     current_link.pop()
                     return True
                 else:

dongtai_engine/filters/utils.py

@@ -0,0 +1,33 @@
+from http.client import parse_headers
+#Request{method=GET, url=http://www.baidu.com/, tag=null}
+
+class JavaObjects:
+
+    def __init__(self, objects_classname, objects_attrs):
+        self.objects_classname = objects_classname
+        self.objects_attrs = objects_attrs
+        for name, value in objects_attrs:
+            setattr(self, name, value)
+
+    def __str__(self):
+        attrs_string = ', '.join(
+            [f'{name}={value}' for name, value in self.objects_attrs])
+        return f"{self.objects_classname}{{{attrs_string}}}"
+
+
+def parse_java_objects(objects_string: str):
+    objects_classname = objects_string[:objects_string.index('{')]
+    objects_attrstring = objects_string[objects_string.index('{'):]
+    objects_attrs = [
+        attr.split('=', 2)
+        for attr in objects_attrstring.strip('{}').split(', ')
+    ]
+    return JavaObjects(objects_classname, objects_attrs)
+
+from io import BytesIO
+from tempfile import TemporaryFile
+def parse_headers_dict_from_bytes(header_bytes: str) -> dict:
+    with TemporaryFile() as fp:
+        fp.write(header_bytes)
+        fp.seek(0)
+        return dict(parse_headers(fp))

dongtai_engine/signals/handlers/vul_handler.py

@@ -50,19 +50,18 @@ def parse_body(body, taint_value):
     except Exception as e:
         return parse_params(body, taint_value)
 
+from dongtai_engine.filters.utils import parse_headers_dict_from_bytes
 
 def parse_header(req_header, taint_value):
     """
     从header头中解析污点的位置
     """
     import base64
-    header_raw = base64.b64decode(req_header).decode('utf-8').split('\n')
-    for header in header_raw:
-        _header_list = header.split(':')
-        _header_name = _header_list[0]
-        _header_value = ':'.join(_header_list[1:])
-        if equals(taint_value, _header_value):
-            return _header_name
+    header_dict = parse_headers_dict_from_bytes(base64.b64decode(req_header))
+    for k, v in header_dict.items():
+        if v == taint_value:
+            return k
+
 
 
 def parse_cookie(req_header, taint_value):

dongtai_engine/tests.py

@@ -37,3 +37,12 @@ def test_benchmark_method_pool_scan(self):
         vul_count_after = IastVulnerabilityModel.objects.all().count()
         assert len(data) == vul_count_after - vul_count_begin
 
+
+class CoreTaskTestCase(AgentTestCase):
+
+    def test_search_method_pool(self):
+        method_pool_id = 500475873
+        method_pool = MethodPool.objects.filter(pk=method_pool_id).first()
+        from dongtai_engine.tasks import search_vul_from_method_pool
+        search_vul_from_method_pool(method_pool.pool_sign, method_pool.agent_id)
+