not redirct websocket request in mitm

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
 
@@ -14,6 +13,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 to_mock_server(flow: http.HTTPFlow):
     raw_url = urlparse(flow.request.url)
     raw_host = raw_url.hostname
@@ -43,6 +50,8 @@ def request(flow: http.HTTPFlow):
     if 'mitm.it' in flow.request.url:
         # Support mitm.it
         return
+    if is_websocket_request(flow):
+        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)
+    })
+
+
 async def make_response_header(proxy_resp_headers: dict, context: LyrebirdProxyContext, data=None):
     response_headers = CIMultiDict()
     for k, v in proxy_resp_headers.items():
@@ -116,6 +126,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:
@@ -198,8 +209,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": {
@@ -211,18 +222,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)