Merge pull request #133 from exexute/main

Fix ( #132 ): add json data pre-check.

core/plugins/strategy_sensitive.py

@@ -90,7 +90,7 @@ def search_id_card_leak(method_pool):
                 save_vul(vul_type='ID Number Leak', method_pool=method_pool, position=key, data=card)
         except Exception as e:
             logger.error(
-                f'check_response_content error, rule name: ID Number Leak, reason: {e}')
+                f'check_response_content error, rule name: ID Number Leak, Method Pool ID: {method_pool.id}, reason: {e}')
 
 
 def check_id_card(id_card):

core/tasks.py

@@ -604,7 +604,7 @@ def export_report():
     logger.info(f'导出报告')
     report = ProjectReport.objects.filter(status=0).first()
     if not report:
-        logger.error("暂无需要导出的报告")
+        logger.info("暂无需要导出的报告")
         return
     try:
         report.status = 1
@@ -648,12 +648,13 @@ def vul_recheck():
             continue
 
         param_name_value = vulnerability['param_name']
-        try:
-            params = json.loads(param_name_value)
-        except JSONDecodeError as e:
-            logger.error(f'污点数据解析出错，原因：{e}')
-            Replay.replay_failed(replay=replay, timestamp=timestamp)
-            continue
+        if param_name_value is not None:
+            try:
+                params = json.loads(param_name_value)
+            except JSONDecodeError as e:
+                logger.error(f'污点数据解析出错，原因：{e}')
+                Replay.replay_failed(replay=replay, timestamp=timestamp)
+                continue
 
         uri = vulnerability['uri']
         param_value = vulnerability['req_params'] if vulnerability['req_params'] else ''

test/core/tasks.py

@@ -17,7 +17,7 @@ def test_search_vul_from_replay_method_pool(self):
         search_vul_from_replay_method_pool(method_id)
 
     def test_search_vul_from_method_pool(self):
-        method_pool_id = 68871
+        method_pool_id = 2473688
         from core.tasks import search_vul_from_method_pool
         search_vul_from_method_pool(method_pool_id)
 
@@ -46,6 +46,58 @@ def test_update_sca(self):
         from core.tasks import update_sca
         update_sca()
 
+    def test_http_header(self):
+        from dongtai.models.agent import IastAgent
+        agents = IastAgent.objects.filter(bind_project_id=1252).values('id')
+        from dongtai.models.agent_method_pool import MethodPool
+        method_pools = MethodPool.objects.filter(agent_id__in=agents).values('req_header_fs')
+
+        from http.server import BaseHTTPRequestHandler
+        class HttpRequest(BaseHTTPRequestHandler):
+            def __init__(self, raw_request):
+                self.body = None
+                self.uri = None
+                self.params = None
+                from io import BytesIO
+                self.rfile = BytesIO(raw_request.encode())
+                self.raw_requestline = self.rfile.readline()
+                self.error_code = self.error_message = None
+                self.parse_request()
+                self.parse_path()
+                self.parse_body()
+                self._cookie_keys = set()
+
+            @property
+            def cookie_keys(self):
+                return self._cookie_keys
+
+            def init_cookie_keys(self):
+                cookies = self.headers.get('cookies').split(';')
+                for cookie in cookies:
+                    self._cookie_keys.add(cookie.strip().split('=')[0])
+
+            def parse_body(self):
+                if self.body is None:
+                    self.body = self.rfile.read().decode('utf-8')
+                return self.body
+
+            def parse_path(self):
+                items = self.path.split('?')
+                self.uri = items[0]
+                self.params = '?'.join(items[1:])
+
+        project_headers = set()
+        project_cookies = set()
+        for method_pool in method_pools:
+            try:
+                request = HttpRequest(method_pool['req_header_fs'])
+                project_headers = project_headers | set(request.headers.keys())
+                # project_cookies = project_cookies | request.cookie_keys
+            except:
+                pass
+        print(project_headers)
+        print(project_cookies)
+
 
 if __name__ == '__main__':
     unittest.main()

vuln/urls.py

@@ -8,11 +8,13 @@
 from django.urls import path
 
 from vuln.views.health import HealthEndPoint
+from vuln.views.proxy import ProxyEndPoint
 from vuln.views.strategy_run import StrategyRunEndPoint
 from vuln.views.sca import ScaEndPoint
 
 urlpatterns = [
     path('run', StrategyRunEndPoint.as_view()),
     path('health', HealthEndPoint.as_view()),
     path('sca', ScaEndPoint.as_view()),
+    path('proxy', ProxyEndPoint.as_view()),
 ]

vuln/views/proxy.py

@@ -0,0 +1,54 @@
+#!/usr/bin/env python
+# -*- coding:utf-8 -*-
+# author: [email protected]
+# datetime: 2022/1/17 下午7:25
+# project: DongTai-engine
+import requests
+from dongtai.endpoint import UserEndPoint, R
+from dongtai.models.agent import IastAgent
+from dongtai.models.agent_method_pool import MethodPool
+
+
+class ProxyEndPoint(UserEndPoint):
+    """
+    引擎注册接口
+    """
+    name = "api-engine-run"
+    description = "引擎运行策略"
+
+    @staticmethod
+    def send_request(method_pool, proxy_addr):
+        url = f'{method_pool["url"]}?{method_pool["req_params"]}' if method_pool['req_params'] else method_pool[
+            'url']
+        header_items = method_pool['req_header_fs'].split('\n')[1:]
+        headers = {}
+        for header_item in header_items:
+            sub_items = header_item.split(':')
+            headers[sub_items[0]] = ':'.join(sub_items[1:])
+        resp = requests.request(
+            method=method_pool["http_method"],
+            data=method_pool['req_data'] if method_pool['req_data'] else '',
+            url=url,
+            headers=headers,
+            proxies={'http': f'http://{proxy_addr}', 'https': f'https://{proxy_addr}'}
+        )
+        print(resp.text)
+
+    def get(self, request):
+        """
+        IAST下载 agent接口 http://localhost:8000/api/engine/proxy?projectId=1195&projectVersionId=992&proxy=http://localhost:5555
+        :param request:
+        :return:
+        """
+        try:
+            project_id = request.query_params.get('projectId')
+            project_version_id = request.query_params.get('projectVersionId')
+            proxy_addr = request.query_params.get('proxy')
+            method_pools = MethodPool.objects.filter(
+                agent__in=IastAgent.objects.filter(bind_project_id=project_id, project_version_id=project_version_id, user=request.user)
+            ).values('url', 'http_method', 'req_params', 'req_data', 'req_header_fs')
+            for method_pool in method_pools:
+                ProxyEndPoint.send_request(method_pool, proxy_addr)
+            return R.success()
+        except Exception as e:
+            return R.failure(msg=f"{e}")