feature/datamanager-mock-data-render-support-escape (#732)

* Escape keyword in jinja2

* update render

* handle 3 kinds of error

* update comment

* update comment

* update case

* update case

lyrebird/mock/blueprints/apis/flow.py

@@ -36,7 +36,6 @@ def get_flow_list_by_filter(filter_obj):
                 url=item['request'].get('url'),
                 scheme=item['request'].get('scheme'),
                 host=item['request'].get('host'),
-                port=item['request'].get('port'),
                 path=item['request'].get('path'),
                 params=unquote(urlencode(item['request']['query'])),
                 method=item['request'].get('method')
@@ -48,6 +47,10 @@ def get_flow_list_by_filter(filter_obj):
             )if item.get('response') and len(item.get('response').keys()) > 1 else {},
             action=item.get('action', [])
         )
+        # Add key `request.port` when port is not default
+        if item['request'].get('port') not in ['443', '80']:
+            info['request']['port'] = item['request'].get('port')
+
         # Add key `proxy_response` into info only if item contains proxy_response
         if item.get('proxy_response'):
             info['proxy_response'] = {

lyrebird/mock/dm/__init__.py

@@ -4,6 +4,7 @@
 import time
 import codecs
 import shutil
+import traceback
 from pathlib import Path
 from urllib.parse import urlparse
 from collections import OrderedDict
@@ -289,21 +290,22 @@ def get_matched_data(self, flow):
 
         return _matched_data
 
-    def _format_respose_data(self, flow):
+    @staticmethod
+    def _format_respose_data(flow):
         # TODO render mock data before response, support more functions
         origin_response_data = flow['response']['data']
 
         try:
             flow['response']['data'] = utils.render_data_with_tojson(flow['response']['data'])
         except Exception:
             flow['response']['data'] = origin_response_data
-            logger.warning(f'Format response string to json error! {flow["request"]["url"]}')
+            logger.warning(f'Format response string to json error! {flow["request"]["url"]}\n {traceback.format_exc()}')
 
         try:
             flow['response']['data'] = utils.render(flow['response']['data'])
         except Exception:
             flow['response']['data'] = origin_response_data
-            logger.warning(f'Format response data error! {flow["request"]["url"]}')
+            logger.warning(f'Format response data error! {flow["request"]["url"]}\n {traceback.format_exc()}')
 
     def _is_match_rule(self, flow, rules):
         return MatchRules.match(flow, rules)

lyrebird/utils.py

@@ -132,6 +132,7 @@ def download(link, input_path):
         for chunck in resp.iter_content():
             f.write(chunck)
 
+
 def render_data_with_tojson(data):
     config_value_tojson_key = config.get('config.value.tojsonKey')
     data_with_tojson = data
@@ -150,6 +151,73 @@ def render_data_with_tojson(data):
         data_with_tojson = re.sub('("{{)'+pattern_group+'(}}")', r'{{\2 | tojson}}', data_with_tojson)
     return data_with_tojson
 
+
+def handle_jinja2_keywords(data, params=None):
+    '''
+    Jinja2 will throw an exception when dealing with unexpected left brackets, but right brackets will not
+    So only left brackets need to be handled
+
+    Handle 3 kinds of unexpected left brackets:
+    1. More than 2 big brackets              `{{{ }}}` `{{# #}}` `{{% %}}`
+    2. Mismatched keyword                    `{{{` `{{#` `{{%`
+    3. Unknown arguments between {{ and }}   `{{unknown}}`
+
+    Convert unexpected brackets into presentable string in Jinja2, such as `var` -> `{{ 'var' }}`
+    The unexpected left brackets above will be convert into:
+    `{{#`          ->  `{{ '{{#' }}`
+    `{{unknown}}`  ->  `{{ '{{' }}unknown{{ '}}' }}`
+    '''
+
+    keywords_pair = {
+        '{{': '}}',
+        '{%': '%}',
+        '{#': '#}'
+    }
+
+    # split by multiple `{` followed by `{` or `#`` or `%`, such as `{{`, `{{{`, `{{{{`, `{#`, `{{#`
+    # EXAMPLE
+    # data = '{{ip}} {{ip {{{ip'
+    # item_list = ['', '{{', 'ip}} ', '{{', 'ip ', '{{{', 'ip']
+    item_list = re.split('({+[{|#|%])', data)
+    if len(item_list) == 1:
+        return data
+
+    left_pattern_index = None
+    for index, item in enumerate(item_list):
+        if index % 2:
+            # 1. Handle more than 2 big brackets
+            if (len(item) > 2) or (item not in keywords_pair):
+                item_list[index] = "{{ '%s' }}" % (item)
+            else:
+                left_pattern_index = index
+            continue
+
+        if left_pattern_index is None:
+            continue
+
+        left_pattern = item_list[left_pattern_index]
+        left_pattern_index = None
+
+        # 2. Handle mismatched keyword
+        right_pattern = keywords_pair[left_pattern]
+        if right_pattern not in item:
+            item_list[index-1] = "{{ '%s' }}" % (item_list[index-1])
+            continue
+
+        # 3. Handle unknown arguments between {{ and }}
+        if left_pattern == '{{':
+            key_n_lefted = item.split('}}', 1)
+            if len(key_n_lefted) != 2:
+                continue
+            key, _ = key_n_lefted
+            if [key for p in params if key.startswith(p)]:
+                continue
+            item_list[index-1] = "{{ '%s' }}" % (item_list[index-1])
+
+    after_data = ''.join(item_list)
+    return after_data
+
+
 def render(data):
     if not isinstance(data, str):
         logger.warning(f'Format error! Expected str, found {type(data)}')
@@ -163,6 +231,8 @@ def render(data):
         'now':  datetime.datetime.now()
     }
 
+    data = handle_jinja2_keywords(data, params)
+
     try:
         template_data = Template(data, keep_trailing_newline=True)
         return template_data.render(params)

tests/test_dm_format.py

@@ -0,0 +1,160 @@
+import re
+import pytest
+from lyrebird import application
+from lyrebird.mock.dm import DataManager
+from lyrebird.config import ConfigManager
+
+
+@pytest.fixture
+def config():
+    _conf = {
+        'ip': '127.0.0.1',
+        'mock.port': 9090,
+        'custom_key': 'custom_value',
+        'config.value.tojsonKey': ['custom.[a-z0-9]{8}(-[a-z0-9]{4}){3}-[a-z0-9]{12}'],
+        'custom.8df051be-4381-41b6-9252-120d9b558bf6': {"custom_key": "custom_value"}
+    }
+    application._cm = ConfigManager()
+    application._cm.config = _conf
+
+
+def test_format_ip(config):
+    flow = {
+        'response': {
+            'data': '{{ip}}'
+        }
+    }
+    DataManager._format_respose_data(flow)
+    data = flow['response']['data']
+    assert data == '127.0.0.1'
+
+
+def test_format_port(config):
+    flow = {
+        'response': {
+            'data': '{{port}}'
+        }
+    }
+    DataManager._format_respose_data(flow)
+    data = flow['response']['data']
+    assert data == '9090'
+
+
+def test_format_config(config):
+    flow = {
+        'response': {
+            'data': "{{config.get('custom_key')}}"
+        }
+    }
+    DataManager._format_respose_data(flow)
+    data = flow['response']['data']
+    assert data == 'custom_value'
+
+
+def test_format_today(config):
+    flow = {
+        'response': {
+            'data': '{{today}}'
+        }
+    }
+    DataManager._format_respose_data(flow)
+    data = flow['response']['data']
+    assert re.fullmatch(r'\d{4}-\d{2}-\d{2}', data)
+
+
+def test_format_today(config):
+    flow = {
+        'response': {
+            'data': "{{now.strftime('%Y-%m-%d %H:%M:%S')}}"
+        }
+    }
+    DataManager._format_respose_data(flow)
+    data = flow['response']['data']
+    assert re.fullmatch(r'\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}', data)
+
+
+def test_format_tojson(config):
+    flow = {
+        'response': {
+            'data': '"keyA":"valueA","keyB":"{{config.get(\'custom.8df051be-4381-41b6-9252-120d9b558bf6\')}}","keyC":"valueC"'
+        }
+    }
+    DataManager._format_respose_data(flow)
+    data = flow['response']['data']
+    assert '"keyA":"valueA"' in data
+    assert '"keyB":{"custom_key": "custom_value"}' in data
+    assert '"keyC":"valueC"' in data
+
+
+
+def test_format_no_param(config):
+    response_data_str = 'ip'
+    flow = {
+        'response': {
+            'data': response_data_str
+        }
+    }
+    DataManager._format_respose_data(flow)
+    data = flow['response']['data']
+    assert data == response_data_str
+
+
+def test_format_empty_parameter(config):
+    response_data_str = '{{}}'
+    flow = {
+        'response': {
+            'data': response_data_str
+        }
+    }
+    DataManager._format_respose_data(flow)
+    data = flow['response']['data']
+    assert data == response_data_str
+
+
+def test_format_with_mismatched_left_keyword(config):
+    mismatched_str = ' ip {{ip {%ip {#ip'
+    response_data_str = '{{ip}}' + mismatched_str
+    flow = {
+        'response': {
+            'data': response_data_str
+        }
+    }
+    DataManager._format_respose_data(flow)
+    data = flow['response']['data']
+    assert data == f'127.0.0.1{mismatched_str}'
+
+
+def test_format_with_mismatched_right_keyword(config):
+    response_data_str = ' ip ip}} }} ip%} ip#}'
+    flow = {
+        'response': {
+            'data': response_data_str
+        }
+    }
+    DataManager._format_respose_data(flow)
+    data = flow['response']['data']
+    assert data == response_data_str
+
+
+def test_format_with_more_than_2_big_brackets(config):
+    response_data_str = '{{{ ip }}} {{{{ ip }}}} {{% loop %}} {{# comment #}}'
+    flow = {
+        'response': {
+            'data': response_data_str
+        }
+    }
+    DataManager._format_respose_data(flow)
+    data = flow['response']['data']
+    assert data == response_data_str
+
+
+def test_format_unknown_parameter(config):
+    response_data_str = '{{unknown}} content {{unknown}} content {{unknown2}}'
+    flow = {
+        'response': {
+            'data': response_data_str
+        }
+    }
+    DataManager._format_respose_data(flow)
+    data = flow['response']['data']
+    assert data == response_data_str