From 93d9594438214a75683df05a050d1af77d41c59a Mon Sep 17 00:00:00 2001 From: noO0oOo0ob <38344038+noO0oOo0ob@users.noreply.github.com> Date: Tue, 24 Dec 2024 14:06:00 +0800 Subject: [PATCH] avoid mitm send lyrebird request to extra mock (#882) Co-authored-by: noO0ob --- lyrebird/mitm/mitm_script.py | 16 ++++++++++++ lyrebird/mitm/proxy_server.py | 3 +++ .../lyrebird_proxy_protocol.py | 4 +++ lyrebird/mock/extra_mock_server/server.py | 25 ++++++++++++++----- 4 files changed, 42 insertions(+), 6 deletions(-) diff --git a/lyrebird/mitm/mitm_script.py b/lyrebird/mitm/mitm_script.py index 5892a660e..cdf913567 100644 --- a/lyrebird/mitm/mitm_script.py +++ b/lyrebird/mitm/mitm_script.py @@ -10,10 +10,23 @@ from urllib.parse import urlparse +SERVER_IP = os.environ.get('SERVER_IP') +MOCK_PORT = int(os.environ.get('MOCK_PORT')) PROXY_PORT = int(os.environ.get('PROXY_PORT')) PROXY_FILTERS = json.loads(os.environ.get('PROXY_FILTERS')) +def check_lyrebird_request(flow: http.HTTPFlow): + parsed_url = urlparse(flow.request.url) + host = parsed_url.hostname + port = parsed_url.port + if host not in ('localhost', '127.0.0.1', SERVER_IP): + return False + if not port or port not in (MOCK_PORT, PROXY_PORT): + return False + return True + + def to_mock_server(flow: http.HTTPFlow): raw_url = urlparse(flow.request.url) raw_host = raw_url.hostname @@ -43,6 +56,9 @@ def request(flow: http.HTTPFlow): if 'mitm.it' in flow.request.url: # Support mitm.it return + if check_lyrebird_request(flow): + # Avoid internal requests + return to_mock_server(flow) diff --git a/lyrebird/mitm/proxy_server.py b/lyrebird/mitm/proxy_server.py index 81663356d..7758b15fc 100644 --- a/lyrebird/mitm/proxy_server.py +++ b/lyrebird/mitm/proxy_server.py @@ -52,6 +52,7 @@ def start_mitmdump(self, queue, config, logger, mitmdump_path): proxy_port = config.get('proxy.port', 4272) mock_port = config.get('mock.port', 9090) extra_mock_port = config.get('extra.mock.port', 9999) + server_ip = config.get('ip', '127.0.0.1') ''' --ignore_hosts: The ignore_hosts option allows you to specify a regex which is matched against a host:port @@ -80,6 +81,8 @@ def start_mitmdump(self, queue, config, logger, mitmdump_path): if ignore_hosts: mitm_arguments += ['--ignore-hosts', ignore_hosts] mitmenv = os.environ + mitmenv['SERVER_IP'] = str(server_ip) + mitmenv['MOCK_PORT'] = str(mock_port) mitmenv['PROXY_PORT'] = str(extra_mock_port) mitmenv['PROXY_FILTERS'] = json.dumps(config.get('proxy.filters', [])) logger.info('HTTP proxy server starting...') diff --git a/lyrebird/mock/extra_mock_server/lyrebird_proxy_protocol.py b/lyrebird/mock/extra_mock_server/lyrebird_proxy_protocol.py index 8f27389ab..91de499b9 100644 --- a/lyrebird/mock/extra_mock_server/lyrebird_proxy_protocol.py +++ b/lyrebird/mock/extra_mock_server/lyrebird_proxy_protocol.py @@ -10,6 +10,9 @@ class UnknownLyrebirdProxyProtocol(Exception): class LyrebirdProxyContext: + + logger = object() + def __init__(self): self.request: web.Request = None self.netloc = None @@ -156,5 +159,6 @@ def parse(cls, request: web.Request, lb_config): for parser in ctx.protocol_parsers: parser(request, lb_config) if ctx.init: + cls.logger.info(f'Hit protocol {parser.__name__}: {ctx.origin_url}') return ctx raise UnknownLyrebirdProxyProtocol diff --git a/lyrebird/mock/extra_mock_server/server.py b/lyrebird/mock/extra_mock_server/server.py index 20962c639..980c7f01d 100644 --- a/lyrebird/mock/extra_mock_server/server.py +++ b/lyrebird/mock/extra_mock_server/server.py @@ -44,6 +44,22 @@ def make_raw_headers_line(request: web.Request): return json.dumps(raw_headers, ensure_ascii=False) +def make_request_headers(context: LyrebirdProxyContext, is_proxy): + headers = {k: v for k, v in context.request.headers.items() if k.lower() not in [ + 'cache-control', 'host', 'transfer-encoding']} + if is_proxy: + if 'Proxy-Raw-Headers' in context.request.headers: + del headers['Proxy-Raw-Headers'] + if 'Lyrebird-Client-Address' in context.request.headers: + del headers['Lyrebird-Client-Address'] + else: + if 'Proxy-Raw-Headers' not in context.request.headers: + headers['Proxy-Raw-Headers'] = make_raw_headers_line(context.request) + if 'Lyrebird-Client-Address' not in context.request.headers: + headers['Lyrebird-Client-Address'] = context.request.remote + return headers + + async def make_response_header(proxy_resp_headers: dict, context: LyrebirdProxyContext, data=None): response_headers = CIMultiDict() for k, v in proxy_resp_headers.items(): @@ -63,12 +79,7 @@ async def make_response_header(proxy_resp_headers: dict, context: LyrebirdProxyC async def send_request(context: LyrebirdProxyContext, target_url): async with client.ClientSession(auto_decompress=False) as session: request: web.Request = context.request - headers = {k: v for k, v in request.headers.items() if k.lower() not in [ - 'cache-control', 'host', 'transfer-encoding']} - if 'Proxy-Raw-Headers' not in request.headers: - headers['Proxy-Raw-Headers'] = make_raw_headers_line(request) - if 'Lyrebird-Client-Address' not in request.headers: - headers['Lyrebird-Client-Address'] = request.remote + headers = make_request_headers(context, target_url==context.origin_url) request_body = None if request.body_exists: request_body = request.content @@ -136,6 +147,8 @@ def init_app(config): log.init(config, logger_queue) logger = log.get_logger() + LyrebirdProxyContext.logger = logger + app = web.Application() app.router.add_route('*', r'/{path:(.*)}', req_handler)