diff --git a/lyrebird/mock/blueprints/apis/flow.py b/lyrebird/mock/blueprints/apis/flow.py index 78c3c29da..e563a4443 100644 --- a/lyrebird/mock/blueprints/apis/flow.py +++ b/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'] = { diff --git a/lyrebird/mock/dm/__init__.py b/lyrebird/mock/dm/__init__.py index b0ef2d28a..5d779d489 100644 --- a/lyrebird/mock/dm/__init__.py +++ b/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,7 +290,8 @@ 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'] @@ -297,13 +299,13 @@ def _format_respose_data(self, flow): 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) diff --git a/lyrebird/utils.py b/lyrebird/utils.py index 485325494..e471baa25 100644 --- a/lyrebird/utils.py +++ b/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) diff --git a/tests/test_dm_format.py b/tests/test_dm_format.py new file mode 100644 index 000000000..1d0fbb935 --- /dev/null +++ b/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