Close ISSUE #92 (Increase the judgment of whether prompt vulnerabilities are enabled) (#93)

Co-authored-by: owefsad <dongzhiyong@secnium.cn>

Author: exexute

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

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,