Fix/narmal vul upload and header vul

.github/workflows/teststate.yml

@@ -120,7 +120,7 @@ jobs:
 
       - name: Django Unit Testing
         run: |
-          mypy --show-error-codes --ignore-missing-imports --disable-error-code var-annotated  --disable-error-code  list-item  --disable-error-code attr-defined --disable-error-code arg-type --disable-error-code assignment --disable-error-code misc .
+          mypy --show-error-codes --ignore-missing-imports  --no-incremental --show-error-codes --check-untyped-defs --disable-error-code var-annotated  --disable-error-code  list-item  --disable-error-code attr-defined --disable-error-code arg-type --disable-error-code assignment --disable-error-code misc --disable-error-code union-attr --disable-error-code index --disable-error-code call-overload  --disable-error-code call-arg --disable-error-code dict-item  --disable-error-code truthy-function --disable-error-code operator --disable-error-code name-defined .
 
   Run-Pep8Check:
     runs-on: ubuntu-latest

dongtai_common/common/utils/__init__.py

@@ -98,6 +98,7 @@ def _noname(function):
 
 
 from rest_framework.authentication import TokenAuthentication, get_authorization_header
+from django.utils.translation import gettext_lazy as _
 
 
 class DepartmentTokenAuthentication(TokenAuthentication):

dongtai_common/models/user.py

@@ -70,7 +70,7 @@ def create_system_user(self,
 
 
 class User(AbstractUser, PermissionsMixin):
-    is_superuser = models.IntegerField(default=0)
+    is_superuser: int = models.IntegerField(default=0)
     phone = models.CharField(max_length=15)
     default_language = models.CharField(max_length=15)
     objects = SaaSUserManager()

dongtai_engine/plugins/strategy_headers.py

@@ -98,7 +98,8 @@ def check_response_header(method_pool):
                      method_pool,
                      position='HTTP Response Header')
     except Exception as e:
-        logger.error("check_response_header failed, reason: " + str(e))
+        logger.error("check_response_header failed, reason: " + str(e),
+                     exc_info=e)
 
 
 def save_vul(vul_type, method_pool, position=None, data=None):
@@ -185,8 +186,8 @@ def save_vul(vul_type, method_pool, position=None, data=None):
             client_ip=method_pool.clent_ip,
             param_name=None,
             method_pool_id=method_pool.id,
-            project_version_id=vul.agent.project_version_id,
-            project_id=vul.agent.bind_project_id,
+            project_version_id=method_pool.agent.project_version_id,
+            project_id=method_pool.agent.bind_project_id,
             language=method_pool.agent.language,
             server_id=method_pool.agent.server_id,
         )

dongtai_protocol/report/handler/narmal_vul_handler.py

@@ -279,7 +279,7 @@ def save(self):
             strategy_id=strategy_id,
             uri=self.http_uri,
             http_method=self.http_method,
-            project_version_id=iast_vul.agent.project_version_id,
+            project_version_id=self.agent.project_version_id,
             pk__lt=iast_vul.id,
         ).delete()
         header_vul = None

dongtai_protocol/utils.py

@@ -102,8 +102,8 @@ def updateossstatus():
                     package_name=package_name),
                 EngineDownloadEndPoint.LOCAL_AGENT_FILE.format(
                     package_name=package_name))
-        downloadstatus = JavaAgentDownload.download_agent(
-        ) and PythonAgentDownload.download_agent()
+        downloadstatus = JavaAgentDownload().download_agent(
+        ) and PythonAgentDownload().download_agent()
         return downloadstatus, None
     except RequestError:
         return False, None

dongtai_web/aggr_vul/aggr_vul_list.py

@@ -39,7 +39,7 @@
 INT_LIMIT: int = 2**64 - 1
 
 
-def convert_cwe(cwe: [List, str]) -> str:
+def convert_cwe(cwe: List | str) -> str:
     if isinstance(cwe, list):
         if len(cwe) > 0:
             return cwe[0].replace("CWE-", "")

dongtai_web/aggr_vul/aggr_vul_summary.py

@@ -286,11 +286,11 @@ def get_annotate_data_es(
             have_poc_count = i['doc_count']
             for k in i['article']['buckets']:
                 if k['key'] == 1:
-                    have_article_count += k['doc_count']
+                    have_article_count += int(k['doc_count'])
         if i['key'] == 0:
             for k in i['article']['buckets']:
                 if k['key'] == 1:
-                    have_article_count += k['doc_count']
+                    have_article_count += int(k['doc_count'])
                 if k['key'] == 0:
                     no_usable_count = k['doc_count']
 

dongtai_web/projecttemplate/update_department_data.py

@@ -10,5 +10,5 @@ def update_department_data():
         else:
             department_dict[department.id] = f"{department_dict[department.parent_id]}-{department.id}"
         department.department_path = department_dict[department.id]
-        department.token = Token.generate_key()
+        department.token = Token().generate_key()
         department.save()

dongtai_web/utils.py

@@ -62,7 +62,7 @@ def assemble_query(condictions: list,
                 }, condictions)), base_query)
 
 
-def assemble_query_2(condictions: dict,
+def assemble_query_2(condictions: list,
                      lookuptype='',
                      base_query=Q(),
                      operator_=operator.or_):

dongtai_web/views/engine_method_pool_search.py

@@ -184,7 +184,7 @@ def post(self, request):
             return R.failure(gettext_lazy("Parameter error"))
         search_fields = dict(
             filter(lambda k: k[0] in fields, request.data.items()))
-        search_fields_ = []
+        search_fields_: list = []
         for k, v in search_fields.items():
             search_fields_.append((k, v))
         search_after_fields = list(
@@ -198,9 +198,9 @@ def post(self, request):
         if 'id' in request.data.keys():
             q = assemble_query(search_after_fields, 'lte', q, operator.and_)
             if search_mode == 1:
-                q = assemble_query(search_fields, 'regex', Q(), operator.or_)
+                q = assemble_query(search_fields_, 'regex', Q(), operator.or_)
             elif search_mode == 2:
-                q = assemble_query_2(search_fields, 'regex', Q(),
+                q = assemble_query_2(search_fields_, 'regex', Q(),
                                      operator.and_)
             if 'id' in request.data.keys():
                 q = q & Q(pk=request.data['id'])

dongtai_web/views/vul_request_replay.py

@@ -91,8 +91,8 @@ def check_method_pool(method_pool_id, user):
 
         auth_agents = RequestReplayEndPoint.get_auth_agents_with_user(user)
         if method_pool_id == -1:
-            method_pool_model = namedtuple('MethodPool', ['id', 'agent'])
-            agent = namedtuple('MethodPool', ['id', 'is_running'])
+            method_pool_model = namedtuple('method_pool_model', ['id', 'agent'])
+            agent = namedtuple('agent', ['id', 'is_running'])
             agent.id = 0
             agent.is_running = 0
             method_pool_model.agent = agent