not redirct websocket request in mitm (#884)

* web not proxy websocket

* add upgrade request report log

* version

---------

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

lyrebird/mitm/mitm_script.py

@@ -3,9 +3,8 @@
 
 Redirect request from proxy server to mock server
 """
-import re
-import json
 import os
+import json
 from mitmproxy import http
 from urllib.parse import urlparse
 
@@ -16,6 +15,14 @@
 PROXY_FILTERS = json.loads(os.environ.get('PROXY_FILTERS'))
 
 
+def is_websocket_request(flow: http.HTTPFlow) -> bool:
+    headers = flow.request.headers
+    return (
+        headers.get('Upgrade', '').lower() == 'websocket' and
+        headers.get('Connection', '').lower() == 'upgrade'
+    )
+
+
 def check_lyrebird_request(flow: http.HTTPFlow):
     parsed_url = urlparse(flow.request.url)
     host = parsed_url.hostname
@@ -59,6 +66,9 @@ def request(flow: http.HTTPFlow):
     if check_lyrebird_request(flow):
         # Avoid internal requests
         return
+    if is_websocket_request(flow):
+        # Avoid Websocket connect requests
+        return
     to_mock_server(flow)
 
 

lyrebird/mock/extra_mock_server/__init__.py

@@ -12,5 +12,4 @@ def run(self, async_obj, config, *args, **kwargs):
         signal.signal(signal.SIGINT, signal.SIG_IGN)
         log_queue = async_obj['logger_queue']
         msg_queue = async_obj['msg_queue']
-        publish_init_status(msg_queue, 'READY')
         serve(msg_queue, config, log_queue, *args, **kwargs)

lyrebird/mock/extra_mock_server/server.py

@@ -15,6 +15,7 @@
 
 logger = None
 logger_queue = None
+message_queue = None
 lb_config = {}
 semaphore = None
 
@@ -44,6 +45,15 @@ def make_raw_headers_line(request: web.Request):
     return json.dumps(raw_headers, ensure_ascii=False)
 
 
+def upgrade_request_report(context: LyrebirdProxyContext):
+    if not context.request.headers.get('upgrade'):
+        return
+    publish_metrics_msg('upgrade_request', {
+        'url': context.origin_url,
+        'headers': str(context.request.headers)
+    })
+
+
 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']}
@@ -127,6 +137,7 @@ async def req_handler(request: web.Request):
     try:
         global lb_config
         proxy_ctx = LyrebirdProxyContext.parse(request, lb_config)
+        upgrade_request_report(proxy_ctx)
         if is_filtered(proxy_ctx):
             # forward to lyrebird
             async with semaphore:
@@ -211,8 +222,8 @@ def _cancel_tasks(
                 }
             )
 
-def publish_init_status(queue, status):
-    queue.put({
+def publish_init_status(status):
+    message_queue.put({
         'type': 'event',
         "channel": "system",
         "content": {
@@ -224,18 +235,34 @@ def publish_init_status(queue, status):
         }
     })
 
+def publish_metrics_msg(action, info):
+    message_queue.put({
+        'type': 'event',
+        "channel": "lyrebird_metrics",
+        "content": {
+            'lyrebird_metrics': {
+                'sender': 'ExtraMockServer',
+                'action': action,
+                'trace_info': str(info)
+                }
+            }
+    })
 
 def serve(msg_queue, config, log_queue, *args, **kwargs):
     global logger_queue
+    global message_queue
     logger_queue = log_queue
+    message_queue = msg_queue
+
+    publish_init_status('READY')
     loop = asyncio.new_event_loop()
     main_task = loop.create_task(_run_app(config))
 
     try:
         asyncio.set_event_loop(loop)
         loop.run_until_complete(main_task)
     except KeyboardInterrupt:
-        publish_init_status(msg_queue, 'ERROR')
+        publish_init_status('ERROR')
     finally:
         _cancel_tasks({main_task}, loop)
         _cancel_tasks(asyncio.all_tasks(loop), loop)

lyrebird/version.py

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