调整代码结构

Author: virusdefender

webhook/app.py

@@ -1,96 +1,22 @@
-import datetime
 import logging
-from datetime import datetime
 
 from flask import Flask, request, redirect
 
 from config import parse_config, get_config
-from executor.executor import dispatch_web_vuln, dispatch_service_vuln, dispatch_statistics
 from executor.registry import init_plugin
-from model.vuln import Statistics, WebVuln, WebParam, WebParamPosition, WebRequest, WebResponse, ServiceVuln
+from views.views import process_web_vuln, process_statistics, process_host_vuln
 
 app = Flask(__name__)
 
-debug = True
 logging.basicConfig(format="[%(levelname)s] %(asctime)s %(name)s %(message)s",
                     datefmt="%Y-%m-%d %H:%M:%S", level=logging.DEBUG)
 
-logger = logging.getLogger("webhook")
-
-
-def process_web_vuln(instance, data):
-    """将 web 漏洞 json 转换为相关 model"""
-    detail = data["detail"]
-    p = detail["param"]
-    if p:
-        param = WebParam(key=p["key"], value=p["value"], position=WebParamPosition(p["position"]))
-    else:
-        param = None
-
-    request = []
-    response = []
-    extra = {}
-
-    for i in range(0, 10):
-        req_key = f"request{i}" if i else "request"
-        resp_key = f"response{i}" if i else "response"
-        req = detail.get(req_key)
-        resp = detail.get(resp_key)
-
-        if req == "" or resp == "":
-            continue
-        if req is None or resp is None:
-            break
-        request.append(WebRequest(raw=req))
-        response.append(WebResponse(raw=resp))
-
-    # 其他的数据可能是自定义的，就单独拿出来
-    not_extra_key = ["request", "response", "param", "payload", "url"]
-    for k, v in detail.items():
-        for item in not_extra_key:
-            if item in k:
-                break
-        else:
-            extra[k] = v
-
-    vuln = WebVuln(create_time=datetime.fromtimestamp(data["create_time"] / 1000), plugin=data["plugin"],
-                   vuln_class=data["vuln_class"],
-                   url=data["target"]["url"], param=param, request=request, response=response, extra=extra)
-    dispatch_web_vuln(instance, vuln)
-
-
-def process_statistics(instance, data):
-    """将统计数据 json 转换为相关 json"""
-    data.pop("type", None)
-    s = Statistics(**data)
-    dispatch_statistics(instance, s)
-
-
-def process_host_vuln(instance, data):
-    """将服务漏洞 json 转换为相关 json"""
-    detail = data["detail"]
-    extra = {}
-
-    not_extra_key = ["host", "port"]
-    for k, v in detail.items():
-        for item in not_extra_key:
-            if item in k:
-                break
-        else:
-            extra[k] = v
-
-    vuln = ServiceVuln(create_time=datetime.fromtimestamp(data["create_time"] / 1000), plugin=data["plugin"],
-                       vuln_class=data["vuln_class"], host=detail["host"], port=detail["port"],
-                       extra=extra)
-    dispatch_service_vuln(instance, vuln)
-
 
 # 比如可以给一个界面管理当前的插件，看到插件的数据等？
 @app.route("/", methods=["GET"])
 def index():
     return redirect("https://xray.cool/xray/#/api/api")
 
-
 @app.route("/webhook", methods=["POST"])
 def webhook():
     token = get_config().server_config.token

webhook/executor/__init__.py

@@ -7,6 +7,7 @@
 @dataclass
 class Statistics:
     """统计数据"""
+    raw_json: dict
     # 发现的 url 数量
     num_found_urls: int
     # 扫描完成的 url 数量
@@ -55,6 +56,7 @@ class WebResponse:
 @dataclass
 class WebVuln:
     """web 漏洞"""
+    raw_json: dict
     # 创建时间
     create_time: datetime
     # 这两个数据内部使用其实是 enum，要不要提供给社区？
@@ -78,6 +80,7 @@ class WebVuln:
 @dataclass
 class ServiceVuln:
     """服务漏洞"""
+    raw_json: dict
     # 同 web 漏洞
     create_time: datetime
     plugin: str

webhook/model/__init__.py

@@ -0,0 +1,78 @@
+import datetime
+from datetime import datetime
+
+from executor.executor import dispatch_web_vuln, dispatch_service_vuln, dispatch_statistics
+from model.vuln import Statistics, WebVuln, WebParam, WebParamPosition, WebRequest, WebResponse, ServiceVuln
+
+
+def process_web_vuln(instance, data):
+    """将 web 漏洞 json 转换为相关 model"""
+    detail = data["detail"]
+    p = detail["param"]
+    if p:
+        param = WebParam(key=p["key"], value=p["value"], position=WebParamPosition(p["position"]))
+    else:
+        param = None
+
+    request = []
+    response = []
+    extra = {}
+
+    for i in range(0, 10):
+        req_key = f"request{i}" if i else "request"
+        resp_key = f"response{i}" if i else "response"
+        req = detail.get(req_key)
+        resp = detail.get(resp_key)
+
+        if req == "" or resp == "":
+            continue
+        if req is None or resp is None:
+            break
+        request.append(WebRequest(raw=req))
+        response.append(WebResponse(raw=resp))
+
+    # 其他的数据可能是自定义的，就单独拿出来
+    not_extra_key = ["request", "response", "param", "payload", "url"]
+    for k, v in detail.items():
+        for item in not_extra_key:
+            if item in k:
+                break
+        else:
+            extra[k] = v
+
+    vuln = WebVuln(create_time=datetime.fromtimestamp(data["create_time"] / 1000), plugin=data["plugin"],
+                   vuln_class=data["vuln_class"],
+                   url=data["target"]["url"], param=param, request=request, response=response, extra=extra,
+                   raw_json=data)
+    dispatch_web_vuln(instance, vuln)
+
+
+def process_statistics(instance, data):
+    """将统计数据 json 转换为相关 json"""
+    s = Statistics(num_found_urls=data["num_found_urls"],
+                   num_scanned_urls=data["num_scanned_urls"],
+                   num_sent_http_requests=data["num_sent_http_requests"],
+                   average_response_time=data["average_response_time"],
+                   ratio_failed_http_requests=data["ratio_failed_http_requests"],
+                   ratio_progress=data["ratio_progress"],
+                   raw_json=data)
+    dispatch_statistics(instance, s)
+
+
+def process_host_vuln(instance, data):
+    """将服务漏洞 json 转换为相关 json"""
+    detail = data["detail"]
+    extra = {}
+
+    not_extra_key = ["host", "port"]
+    for k, v in detail.items():
+        for item in not_extra_key:
+            if item in k:
+                break
+        else:
+            extra[k] = v
+
+    vuln = ServiceVuln(create_time=datetime.fromtimestamp(data["create_time"] / 1000), plugin=data["plugin"],
+                       vuln_class=data["vuln_class"], host=detail["host"], port=detail["port"],
+                       extra=extra, raw_json=data)
+    dispatch_service_vuln(instance, vuln)