From b5aa5129afe4fd7928711a4e12981e1c0b1e3780 Mon Sep 17 00:00:00 2001 From: owefsad Date: Tue, 15 Feb 2022 18:56:14 +0800 Subject: [PATCH] Fix ( #132 ): add json data pre-check. --- core/plugins/strategy_sensitive.py | 2 +- core/tasks.py | 15 +++++---- test/core/tasks.py | 54 +++++++++++++++++++++++++++++- vuln/urls.py | 2 ++ vuln/views/proxy.py | 54 ++++++++++++++++++++++++++++++ 5 files changed, 118 insertions(+), 9 deletions(-) create mode 100644 vuln/views/proxy.py diff --git a/core/plugins/strategy_sensitive.py b/core/plugins/strategy_sensitive.py index 53f9189..41b63a8 100644 --- a/core/plugins/strategy_sensitive.py +++ b/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): diff --git a/core/tasks.py b/core/tasks.py index 38599e9..de0de53 100644 --- a/core/tasks.py +++ b/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 '' diff --git a/test/core/tasks.py b/test/core/tasks.py index bc8d220..1e2f4ee 100644 --- a/test/core/tasks.py +++ b/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() diff --git a/vuln/urls.py b/vuln/urls.py index b5b720c..20accc3 100644 --- a/vuln/urls.py +++ b/vuln/urls.py @@ -8,6 +8,7 @@ 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 @@ -15,4 +16,5 @@ path('run', StrategyRunEndPoint.as_view()), path('health', HealthEndPoint.as_view()), path('sca', ScaEndPoint.as_view()), + path('proxy', ProxyEndPoint.as_view()), ] diff --git a/vuln/views/proxy.py b/vuln/views/proxy.py new file mode 100644 index 0000000..b73bf59 --- /dev/null +++ b/vuln/views/proxy.py @@ -0,0 +1,54 @@ +#!/usr/bin/env python +# -*- coding:utf-8 -*- +# author: owefsad@huoxian.cn +# 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}")