diff --git a/sysom_api/apps/vul/models.py b/sysom_api/apps/vul/models.py index 6b2df7f973f014fd77ecc36e04d4d0610481998e..1c175012d68499561ce00bdefb30aaf21a2f0111 100644 --- a/sysom_api/apps/vul/models.py +++ b/sysom_api/apps/vul/models.py @@ -41,18 +41,6 @@ class VulAddrModel(models.Model): def __str__(self): return f'vul addres: {self.url}' - def get_req_arg(self): - headers = self.headers - if "User-Agent" not in headers: - headers[ - "User-Agent"] = "Mozilla/5.0 (X11; Linux x86_64) Chrome/99.0.4844.51" - - if self.authorization_type.lower() == "basic" and self.authorization_body: - auth = self.authorization_body - else: - auth = {} - return self.url, self.get_method_display(), headers, self.params, self.body, auth - class VulBaseModel(BaseModel): cve_id = models.CharField(max_length=100) diff --git a/sysom_api/apps/vul/serializer.py b/sysom_api/apps/vul/serializer.py index a287b3d6916274d50b604d9ef732cd0a8b441a5a..0ce9e7b93d0e51c9b7b5ca77397538db93cb7735 100644 --- a/sysom_api/apps/vul/serializer.py +++ b/sysom_api/apps/vul/serializer.py @@ -1,6 +1,6 @@ # -*- encoding: utf-8 -*- """ -@File : urls.py +@File : serializer.py @Time : 2022/4/8 下午1:49 @Author : weidongkl @Email : weidong@uniontech.com @@ -12,11 +12,13 @@ from apps.vul.models import VulAddrModel class VulAddrListSerializer(serializers.ModelSerializer): method_display = serializers.SerializerMethodField() + headers = serializers.SerializerMethodField() + description = serializers.SerializerMethodField() class Meta: model = VulAddrModel fields = ["id", "name", "description", "method", "method_display", "url", "headers", "params", "body", - "authorization_type", "parser", "status" + "authorization_type", "parser", "status", "is_edited" ] def get_method_display(self, attr: VulAddrModel) -> int: @@ -25,6 +27,18 @@ class VulAddrListSerializer(serializers.ModelSerializer): def get_description(self, attr: VulAddrModel) -> str: return attr.description or '暂未填写' + def get_headers(self, attr: VulAddrModel) -> str: + if attr.is_edited: + return attr.headers + else: + shadow_string = "x" * 12 + shadow_fields = ["token", "authorization"] + display_headers = attr.headers.copy() + for k, v in attr.headers.items(): + if k.lower() in shadow_fields: + display_headers[k] = shadow_string + return display_headers + class VulAddrModifySerializer(serializers.ModelSerializer): authorization_body = serializers.JSONField(required=False) diff --git a/sysom_api/apps/vul/views.py b/sysom_api/apps/vul/views.py index d35e6ae01c68d9b31fcc0492a51b05182cfccb79..d574a30174c5dccb975473b46f8bba367d185a58 100644 --- a/sysom_api/apps/vul/views.py +++ b/sysom_api/apps/vul/views.py @@ -332,16 +332,35 @@ class VulAddrViewSet(viewsets.ModelViewSet): super().update(request, *args, **kwargs) return success(result={}, message="修改成功") - @action(detail=True, methods=['get']) + @action(detail=False, methods=['post']) def test_connect(self, request, *args, **kwargs): - vul = self.get_object() - url, method, headers, params, payload, auth = vul.get_req_arg() + body = request.data + url, method, headers, params, payload, auth = self.get_req_arg(body) req = requests.Request(method, url, headers=headers, data=payload, params=params, auth=auth) prepped = req.prepare() data = {"request": self.get_req_struct(prepped), "status": self.get_resp_result(prepped)} return success(result=data, message="") + @staticmethod + def get_req_arg(body): + headers = body.get("headers") + if "User-Agent" not in headers: + headers[ + "User-Agent"] = "Mozilla/5.0 (X11; Linux x86_64) Chrome/99.0.4844.51" + + if body.get("authorization_type").lower() == "basic" and body.get("authorization_body"): + auth = body.get("authorization_body") + else: + auth = {} + + for i in VulAddrModel.REQUEST_METHOD_CHOICES: + if i[0] == body.get("method"): + method = i[1] + break + + return body.get("url"), method, headers, body.get("params"), body.get("body"), auth + @staticmethod def get_req_struct(req): req_struct = '{}\n\n{}\n\n{}'.format( @@ -361,7 +380,7 @@ class VulAddrViewSet(viewsets.ModelViewSet): msg = f"Status Code: {resp_status} OK" else: msg = f"Status Code: {resp_status} ERROR" + return msg except Exception as e: msg = f"Status Code: ERROR({e})" - finally: return msg