merge master

lyrebird/checker/__init__.py

@@ -264,7 +264,9 @@ def _load_checker(self):
 
         self.update = True
         self._checker_detection()
+
         global scripts_tmp_storage
+        self._funcs_map.clear()
         for func_type, func_list in scripts_tmp_storage.items():
             self._funcs_map[func_type] = [f for f in func_list]
         scripts_tmp_storage = {}

lyrebird/mitm/mitm_script.py

@@ -9,6 +9,8 @@
 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'))
 
@@ -21,6 +23,17 @@ def is_websocket_request(flow: http.HTTPFlow) -> bool:
     )
 
 
+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
@@ -50,7 +63,11 @@ 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
     if is_websocket_request(flow):
+        # Avoid Websocket connect 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

@@ -54,6 +54,22 @@ def upgrade_request_report(context: LyrebirdProxyContext):
     })
 
 
+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():
@@ -73,12 +89,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
@@ -147,6 +158,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)
 

lyrebird/mock/handlers/flow_editor_handler.py

@@ -50,7 +50,7 @@ def on_response_handler(self, handler_context):
             handler_context.flow['response']['code'] = 200
         if not handler_context.flow['response'].get('headers'):
             handler_context.flow['response']['headers'] = {}
-        if not handler_context.flow['response'].get('data'):
+        if 'data' not in handler_context.flow['response']:
             if handler_context.response:
                 handler_context.update_response_data2flow()
             else:

lyrebird/mock/handlers/handler_context.py

@@ -309,8 +309,10 @@ def update_client_resp_time(self):
         resp_data = self.flow['response'].get('data', '')
         if isinstance(resp_data, str):
             self.flow['size'] = len(resp_data.encode())
-        else:
+        elif resp_data:
             self.flow['size'] = len(resp_data)
+        else:
+            self.flow['size'] = 0
 
         self.flow['duration'] = self.server_resp_time - self.client_req_time