avoid mitm send lyrebird request to extra mock (#882)

Co-authored-by: noO0ob <[email protected]>

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)
 
 

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...')

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

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)