add param for on_request (#824)

* add param for on_request

* support plugin param

* change name

* review

* add plugin upstream param

lyrebird/checker/__init__.py

@@ -1,5 +1,6 @@
 import codecs
 import shutil
+import functools
 from pathlib import Path
 from importlib import machinery
 import lyrebird
@@ -324,5 +325,23 @@ def write(self, data):
     def json(self):
         return {k: self.__dict__[k] for k in self.__dict__ if not k.startswith('_')}
 
+
+class DecoratorUtils:
+    @staticmethod
+    def modifiy_request_body_decorator(func, modifiy_request_body):
+        # When the request modifier modifies only headers or urls, 
+        # ensure that the Origin request body switch is still in effect after the request modifier is triggered
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            result = func(*args, **kwargs)
+            if isinstance(args, (list, tuple)) and len(args) > 0 and isinstance(args[0], dict):
+                if 'keep_origin_request_body' in args[0]:
+                    # When multiple request modifiers are triggered, the original request data is not used as long as one modifies the requestBody
+                    args[0]['keep_origin_request_body'] = args[0]['keep_origin_request_body'] and not modifiy_request_body
+                else:
+                    args[0]['keep_origin_request_body'] = not modifiy_request_body
+            return result
+        return wrapper
+
 class CheckerIllegal(Exception):
     pass

lyrebird/checker/on_request.py

@@ -4,7 +4,7 @@
 
 class OnRequestHandler:
 
-    def __call__(self, rules=None, rank=0, *args, **kw):
+    def __call__(self, rules=None, rank=0, modifiy_request_body=True, *args, **kw):
         def func(origin_func):
             func_type = checker.TYPE_ON_REQUEST
             if not checker.scripts_tmp_storage.get(func_type):
@@ -13,18 +13,21 @@ def func(origin_func):
                 'name': origin_func.__name__,
                 'func': origin_func,
                 'rules': rules,
-                'rank': rank if isinstance(rank, (int, float)) else 0
+                'rank': rank if isinstance(rank, (int, float)) else 0,
+                'modifiy_request_body': modifiy_request_body
             })
             return origin_func
         return func
 
     @staticmethod
     def register(func_info):
+        func_info['func'] = checker.DecoratorUtils.modifiy_request_body_decorator(func_info['func'], func_info['modifiy_request_body'])
         application.on_request.append(func_info)
 
     @staticmethod
     def unregister(func_info):
         if func_info in application.on_request:
             application.on_request.remove(func_info)
 
+
 on_request = OnRequestHandler()

lyrebird/checker/on_request_upstream.py

@@ -4,7 +4,7 @@
 
 class OnRequestUpstreamHandler:
 
-    def __call__(self, rules=None, rank=0, *args, **kw):
+    def __call__(self, rules=None, rank=0, modifiy_request_body=True, *args, **kw):
         def func(origin_func):
             func_type = checker.TYPE_ON_REQUEST_UPSTREAM
             if not checker.scripts_tmp_storage.get(func_type):
@@ -13,13 +13,15 @@ def func(origin_func):
                 'name': origin_func.__name__,
                 'func': origin_func,
                 'rules': rules,
-                'rank': rank if isinstance(rank, (int, float)) else 0
+                'rank': rank if isinstance(rank, (int, float)) else 0,
+                'modifiy_request_body': modifiy_request_body
             })
             return origin_func
         return func
 
     @staticmethod
     def register(func_info):
+        func_info['func'] = checker.DecoratorUtils.modifiy_request_body_decorator(func_info['func'], func_info['modifiy_request_body'])
         application.on_request_upstream.append(func_info)
 
     @staticmethod

lyrebird/mock/handlers/handler_context.py

@@ -184,7 +184,7 @@ def set_response_source_proxy(self):
         self.response_source = 'proxy'
 
     def get_request_body(self, in_request_handler=True):
-        if self.is_request_edited:
+        if self.is_request_edited and not self.flow.get('keep_origin_request_body', False):
             # TODO Repeated calls, remove it
             self.flow['request']['headers'] = HeadersHelper.flow2origin(self.flow['request'], chain=self.request_chain)
 
@@ -193,10 +193,10 @@ def get_request_body(self, in_request_handler=True):
             if in_request_handler:
                 _data = self.request.data or self.request.form or None
             # When origin_request is not saved, the original data cannot be obtained when diff-mode is enabled.
-            elif self.request_origin_data:
-                _data = self.request_origin_data
             else:
-                _data = DataHelper.flow2origin(self.flow['request'])
+                _data = self.request_origin_data
+            if self.is_request_edited:
+                logger.info(f'requestBody uses the original data. Please make sure that the modifier does not modify the requestBody in request: {self.flow["request"]["url"]}')
         return _data
 
     def get_request_headers(self):

lyrebird/plugins/plugin_manager.py

@@ -1,6 +1,10 @@
 from lyrebird.base_server import StaticServer
 from . import plugin_loader
 from lyrebird import application
+from lyrebird.checker.on_request import OnRequestHandler
+from lyrebird.checker.on_response import OnResponseHandler
+from lyrebird.checker.on_request_upstream import OnRequestUpstreamHandler
+from lyrebird.checker.on_response_upstream import OnResponseUpstreamHandler
 from flask import Blueprint, send_file, request
 from lyrebird.log import get_logger
 from types import FunctionType
@@ -84,16 +88,17 @@ def print_plugin_api(response):
 
             # Subscribe handler on request
             for handler in plugin.manifest.on_request:
-                application.on_request.append({
+                OnRequestHandler.register({
                     'name': handler[0],
                     'func': handler[1],
                     'rules': handler[2] if len(handler) > 2 else None,
-                    'rank': handler[3] if len(handler) > 3 and isinstance(handler[3], (int, float)) else 0
+                    'rank': handler[3] if len(handler) > 3 and isinstance(handler[3], (int, float)) else 0,
+                    'modifiy_request_body': handler[4] if len(handler) > 4 else True
                 })
 
             # Subscribe handler on response
             for handler in plugin.manifest.on_response:
-                application.on_response.append({
+                OnResponseHandler.register({
                     'name': handler[0],
                     'func': handler[1],
                     'rules': handler[2] if len(handler) > 2 else None,
@@ -102,16 +107,17 @@ def print_plugin_api(response):
 
             # Subscribe handler on proxy request
             for handler in plugin.manifest.on_request_upstream:
-                application.on_request_upstream.append({
+                OnRequestUpstreamHandler.register({
                     'name': handler[0],
                     'func': handler[1],
                     'rules': handler[2] if len(handler) > 2 else None,
-                    'rank': handler[3] if len(handler) > 3 and isinstance(handler[3], (int, float)) else 0
+                    'rank': handler[3] if len(handler) > 3 and isinstance(handler[3], (int, float)) else 0,
+                    'modifiy_request_body': handler[4] if len(handler) > 4 else True
                 })
 
             # Subscribe handler on proxy response
             for handler in plugin.manifest.on_response_upstream:
-                application.on_response_upstream.append({
+                OnResponseUpstreamHandler.register({
                     'name': handler[0],
                     'func': handler[1],
                     'rules': handler[2] if len(handler) > 2 else None,

lyrebird/version.py

@@ -1,3 +1,3 @@
-IVERSION = (2, 24, 1)
+IVERSION = (2, 25, 0)
 VERSION = ".".join(str(i) for i in IVERSION)
 LYREBIRD = "Lyrebird " + VERSION