From a8619c52f2338436d19133c29e895dedb84cabcb Mon Sep 17 00:00:00 2001 From: noO0oOo0ob <38344038+noO0oOo0ob@users.noreply.github.com> Date: Mon, 5 Feb 2024 14:20:24 +0800 Subject: [PATCH] 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 | 19 +++++++++++++++++++ lyrebird/checker/on_request.py | 7 +++++-- lyrebird/checker/on_request_upstream.py | 6 ++++-- lyrebird/mock/handlers/handler_context.py | 8 ++++---- lyrebird/plugins/plugin_manager.py | 18 ++++++++++++------ lyrebird/version.py | 2 +- 6 files changed, 45 insertions(+), 15 deletions(-) diff --git a/lyrebird/checker/__init__.py b/lyrebird/checker/__init__.py index db9b467bc..adda6a7c8 100644 --- a/lyrebird/checker/__init__.py +++ b/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 diff --git a/lyrebird/checker/on_request.py b/lyrebird/checker/on_request.py index b73bbf406..3a71d9d4d 100644 --- a/lyrebird/checker/on_request.py +++ b/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,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.append(func_info) @staticmethod @@ -27,4 +29,5 @@ def unregister(func_info): if func_info in application.on_request: application.on_request.remove(func_info) + on_request = OnRequestHandler() diff --git a/lyrebird/checker/on_request_upstream.py b/lyrebird/checker/on_request_upstream.py index c506364cf..2cef2f5f2 100644 --- a/lyrebird/checker/on_request_upstream.py +++ b/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 diff --git a/lyrebird/mock/handlers/handler_context.py b/lyrebird/mock/handlers/handler_context.py index a5aaf8136..17198871c 100644 --- a/lyrebird/mock/handlers/handler_context.py +++ b/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): diff --git a/lyrebird/plugins/plugin_manager.py b/lyrebird/plugins/plugin_manager.py index e68f352b8..04b3c8af7 100644 --- a/lyrebird/plugins/plugin_manager.py +++ b/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, diff --git a/lyrebird/version.py b/lyrebird/version.py index e6e865b0d..3354da2a1 100644 --- a/lyrebird/version.py +++ b/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