From b2478556503d04827a79a29b3aa906d46ab9167a Mon Sep 17 00:00:00 2001 From: owefsad Date: Tue, 23 Nov 2021 23:29:59 +0800 Subject: [PATCH 1/2] Close ISSUE #92 (Increase the judgment of whether prompt vulnerabilities are enabled) --- core/plugins/__init__.py | 27 +++++++++++++++++++++++++++ core/plugins/strategy_headers.py | 4 ++++ 2 files changed, 31 insertions(+) diff --git a/core/plugins/__init__.py b/core/plugins/__init__.py index c387e38..8721e91 100644 --- a/core/plugins/__init__.py +++ b/core/plugins/__init__.py @@ -3,3 +3,30 @@ # author: owefsad@huoxian.cn # datetime: 2021/10/22 下午2:26 # project: DongTai-engine + +from dongtai.models.project import IastProject +from dongtai.models.strategy import IastStrategyModel +from dongtai.utils import const + + +def is_strategy_enable(vul_type, method_pool): + try: + vul_strategy = IastStrategyModel.objects.filter( + vul_type=vul_type, + state=const.STRATEGY_ENABLE, + user_id__in=(1, method_pool.agent.user.id) + ).first() + if vul_strategy is None: + return False + project_id = method_pool.agent.bind_project_id + project = IastProject.objects.filter(id=project_id).first() + if project is None: + return False + strategy_ids = project.scan.content + if strategy_ids is None: + return False + if str(vul_strategy.id) in strategy_ids.split(','): + return True + return False + except Exception as e: + return False diff --git a/core/plugins/strategy_headers.py b/core/plugins/strategy_headers.py index fcac0fc..da9f339 100644 --- a/core/plugins/strategy_headers.py +++ b/core/plugins/strategy_headers.py @@ -13,6 +13,8 @@ from dongtai.models.vulnerablity import IastVulnerabilityModel from dongtai.utils import const +from core.plugins import is_strategy_enable + class FakeSocket(): def __init__(self, response_str): @@ -82,6 +84,8 @@ def check_response_header(method_pool): def save_vul(vul_type, method_pool, position=None, data=None): + if is_strategy_enable(vul_type, method_pool) is False: + return None vul_strategy = IastStrategyModel.objects.filter( vul_type=vul_type, state=const.STRATEGY_ENABLE, From e1bf53f41f1a139a09694f3876664a19e2869c4f Mon Sep 17 00:00:00 2001 From: owefsad Date: Tue, 30 Nov 2021 15:00:51 +0800 Subject: [PATCH 2/2] Close ISSUE #94 ( DongTai-Engine ) --- core/plugins/strategy_headers.py | 2 ++ core/tasks.py | 2 +- signals/handlers/vul_handler.py | 4 ++++ test/core/tasks.py | 2 +- 4 files changed, 8 insertions(+), 2 deletions(-) diff --git a/core/plugins/strategy_headers.py b/core/plugins/strategy_headers.py index da9f339..da9ea48 100644 --- a/core/plugins/strategy_headers.py +++ b/core/plugins/strategy_headers.py @@ -9,6 +9,7 @@ from celery.apps.worker import logger from django.db.models import Q +from dongtai.models.project import IastProject from dongtai.models.strategy import IastStrategyModel from dongtai.models.vulnerablity import IastVulnerabilityModel from dongtai.utils import const @@ -101,6 +102,7 @@ def save_vul(vul_type, method_pool, position=None, data=None): method_pool_id=method_pool.id ).first() timestamp = int(time.time()) + IastProject.objects.filter(id=method_pool.agent.bind_project_id).update(latest_time=timestamp) if vul: vul.req_header = method_pool.req_header vul.req_params = method_pool.req_params diff --git a/core/tasks.py b/core/tasks.py index fb18bdb..5327741 100644 --- a/core/tasks.py +++ b/core/tasks.py @@ -117,7 +117,6 @@ def search_and_save_vul(engine, method_pool_model, method_pool, strategy): taint_value=taint_value ) else: - # 更新漏洞状态为已忽略/误报 try: if isinstance(method_pool_model, MethodPool): return @@ -140,6 +139,7 @@ def search_and_save_vul(engine, method_pool_model, method_pool, strategy): verify_time=timestamp, update_time=timestamp ) + IastProject.objects.filter(id=method_pool.agent.bind_project_id).update(latest_time=timestamp) except Exception as e: logger.info(f'漏洞数据处理出错，原因：{e}') diff --git a/signals/handlers/vul_handler.py b/signals/handlers/vul_handler.py index c320f40..d9a344f 100644 --- a/signals/handlers/vul_handler.py +++ b/signals/handlers/vul_handler.py @@ -9,6 +9,7 @@ import requests from celery.apps.worker import logger from django.dispatch import receiver +from dongtai.models.project import IastProject from dongtai.models.replay_queue import IastReplayQueue from dongtai.models.notify_config import IastNotifyConfig @@ -239,6 +240,7 @@ def save_vul(vul_meta, vul_level, strategy_id, vul_stack, top_stack, bottom_stac agent=vul_meta.agent, method_pool_id=vul_meta.id ).first() + IastProject.objects.filter(id=vul_meta.agent.bind_project_id).update(latest_time=timestamp) if vul: vul.req_header = vul_meta.req_header vul.req_params = vul_meta.req_params @@ -327,6 +329,8 @@ def handler_replay_vul(vul_meta, vul_level, strategy_id, vul_stack, top_stack, b vul.latest_time = timestamp vul.save(update_fields=['status_id', 'latest_time']) + IastProject.objects.filter(id=vul_meta.agent.bind_project_id).update(latest_time=timestamp) + IastReplayQueue.objects.filter(id=kwargs['replay_id']).update( state=const.SOLVED, result=const.RECHECK_TRUE, diff --git a/test/core/tasks.py b/test/core/tasks.py index a2823da..bc8d220 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 = 66235 + method_pool_id = 68871 from core.tasks import search_vul_from_method_pool search_vul_from_method_pool(method_pool_id)