fix-edge-case

iast/views/agent_download.py

@@ -12,7 +12,7 @@
 from django.utils.translation import gettext_lazy as _
 from dongtai.models.profile import IastProfile
 from iast.utils import get_openapi
-
+from requests.exceptions import ConnectionError
 
 logger = logging.getLogger('dongtai-webapi')
 
@@ -60,12 +60,18 @@ def get(self, request):
         token, success = Token.objects.values('key').get_or_create(user=request.user)
         AGENT_SERVER_PROXY={'HOST':''}
         AGENT_SERVER_PROXY['HOST'] = get_openapi()
-        resp = requests.get(
+        try:
+            resp = requests.get(
             url=f'{AGENT_SERVER_PROXY["HOST"]}/api/v1/agent/download?url={base_url}&language={language}&projectName={project_name}',
             headers={
                 'Authorization': f'Token {token["key"]}'
             })
-
+        except ConnectionError as e:
+            return R.failure(msg='conncet error,please check config.ini')
+        except Exception as e:
+            logger.error(e)
+            return R.failure(msg='download error,please check deployment')
+
         response = self.res_by_language(language, token, resp)
 
         return response

iast/views/strategys.py

@@ -22,6 +22,8 @@
 from rest_framework.serializers import ValidationError
 from dongtai.permissions import TalentAdminPermission
 from rest_framework.decorators import permission_classes
+from dongtai.models.vul_level import IastVulLevel
+
 class _StrategyResponseDataStrategySerializer(serializers.Serializer):
     id = serializers.CharField(help_text=_('The id of agent'))
     vul_name = serializers.CharField(help_text=_('The name of the vulnerability type targeted by the strategy'))
@@ -51,14 +53,20 @@ class StrategyCreateSerializer(serializers.Serializer):
         _("Suggestions for repairing vulnerabilities corresponding to the strategy"
           ))
 
+    def validate_level_id(self, value):
+        if not IastVulLevel.objects.filter(pk=value).exists():
+            raise serializers.ValidationError("this vul level not exist")
+        return value
+
 
 _ResponseSerializer = get_response_serializer(
     data_serializer=_StrategyResponseDataStrategySerializer(many=True), )
 
 
 class _StrategyArgsSerializer(serializers.Serializer):
-    page_size = serializers.IntegerField(default=None,help_text=_('Number per page'))
-    page = serializers.IntegerField(default=None,help_text=_('Page index'))
+    page_size = serializers.IntegerField(default=None,
+                                         help_text=_('Number per page'))
+    page = serializers.IntegerField(default=None, help_text=_('Page index'))
     name = serializers.CharField(
         default=None,
         help_text=_(
@@ -175,15 +183,28 @@ def post(self, request):
                 pass
         except ValidationError as e:
             return R.failure(data=e.detail)
+        print(ser.validated_data)
         strategy = IastStrategyModel.objects.create(**ser.validated_data,
-                user=request.user,dt=time.time())
+                                                    user=request.user,
+                                                    dt=time.time())
+        strategy.save()
         for language in IastProgramLanguage.objects.all():
-            HookType.objects.create(type=3,name=ser.validated_data['vul_name'],
-                value=ser.validated_data['vul_type'],enable=1,
-                create_time=time.time(),update_time=time.time(),
-                created_by=request.user.id,language=language,vul_strategy=strategy)
-            HookType.objects.create(type=4,name=ser.validated_data['vul_name'],
-                value=ser.validated_data['vul_type'],enable=1,
-                create_time=time.time(),update_time=time.time(),
-                created_by=request.user.id,language=language,vul_strategy=strategy)
+            HookType.objects.create(type=3,
+                                    name=ser.validated_data['vul_name'],
+                                    value=ser.validated_data['vul_type'],
+                                    enable=1,
+                                    create_time=time.time(),
+                                    update_time=time.time(),
+                                    created_by=request.user.id,
+                                    language=language,
+                                    vul_strategy=strategy)
+            HookType.objects.create(type=4,
+                                    name=ser.validated_data['vul_name'],
+                                    value=ser.validated_data['vul_type'],
+                                    enable=1,
+                                    create_time=time.time(),
+                                    update_time=time.time(),
+                                    created_by=request.user.id,
+                                    language=language,
+                                    vul_strategy=strategy)
         return R.success(data=StrategySerializer(strategy).data)

iast/views/user_login.py

@@ -26,26 +26,30 @@ class UserLogin(UserEndPoint):
         'captcha': ""
     })
     def post(self, request):
-        captcha_hash_key = request.data["captcha_hash_key"]
-        captcha = request.data["captcha"]
-        if captcha_hash_key and captcha:
-            captcha_obj = CaptchaStore.objects.get(hashkey=captcha_hash_key)
-            if int(captcha_obj.expiration.timestamp()) < int(time.time()):
-                return R.failure(status=203, msg=_('Captcha timed out'))
-            if captcha_obj.response == captcha.lower():
-                username = request.data["username"]
-                password = request.data["password"]
-                user = authenticate(username=username, password=password)
-                if user is not None and user.is_active:
-                    login(request, user)
-                    return R.success(
-                        msg=_('Login successful'),
-                        data={'default_language': user.default_language})
+        try:
+            captcha_hash_key = request.data["captcha_hash_key"]
+            captcha = request.data["captcha"]
+            if captcha_hash_key and captcha:
+                captcha_obj = CaptchaStore.objects.get(hashkey=captcha_hash_key)
+                if int(captcha_obj.expiration.timestamp()) < int(time.time()):
+                    return R.failure(status=203, msg=_('Captcha timed out'))
+                if captcha_obj.response == captcha.lower():
+                    username = request.data["username"]
+                    password = request.data["password"]
+                    user = authenticate(username=username, password=password)
+                    if user is not None and user.is_active:
+                        login(request, user)
+                        return R.success(
+                            msg=_('Login successful'),
+                            data={'default_language': user.default_language})
+                    else:
+                        logger.warn(
+                            f"user [{username}] login failure, rease: {'user not exist' if user is None else 'user is disable'}")
+                        return R.failure(status=202, msg=_('Login failed'))
                 else:
-                    logger.warn(
-                        f"user [{username}] login failure, rease: {'user not exist' if user is None else 'user is disable'}")
-                    return R.failure(status=202, msg=_('Login failed'))
+                    return R.failure(status=203, msg=_('Verification code error'))
             else:
-                return R.failure(status=203, msg=_('Verification code error'))
-        else:
-            return R.failure(status=204, msg=_('verification code should not be empty'))
+                return R.failure(status=204, msg=_('verification code should not be empty'))
+        except Exception as e:
+            logger.error(e)
+            return R.failure(status=202, msg=_('Login failed'))

iast/views/user_passwrd.py

@@ -4,29 +4,35 @@
 
 # software: PyCharm
 # project: lingzhi-webapi
+import logging
 from django.contrib.auth import authenticate
 
 from dongtai.endpoint import R
 from dongtai.endpoint import UserEndPoint
 from django.utils.translation import gettext_lazy as _
 
+logger = logging.getLogger("dongtai-webapi")
 
 class UserPassword(UserEndPoint):
     name = "api-v1-user-password"
     description = _("Change Password")
 
     def post(self, request):
         user = request.user
-
-        if not request.data['old_password'] or not request.data['new_password']:
-            return R.failure(msg=_('Password should not be empty'))
-        else:
-            user_check = authenticate(username=user.username, password=request.data['old_password'])
-            if user_check is not None and user_check.is_active:
-                password = request.data['new_password']
-
-                user.set_password(password)
-                user.save(update_fields=['password'])
-                return R.success(msg=_('Password has been changed successfully'))
+        try:        
+            if not request.data['old_password'] or not request.data['new_password']:
+                return R.failure(msg=_('Password should not be empty'))
             else:
-                return R.failure(msg=_('Incorrect old password'))
+                user_check = authenticate(username=user.username, password=request.data['old_password'])
+                if user_check is not None and user_check.is_active:
+                    password = request.data['new_password']
+
+                    user.set_password(password)
+                    user.save(update_fields=['password'])
+                    return R.success(msg=_('Password has been changed successfully'))
+                else:
+                    return R.failure(msg=_('Incorrect old password'))
+        except Exception as e:
+            logger.error(e)
+            return R.failure(msg=_('Incorrect'))
+