stoobly-agent 1.3.0__py3-none-any.whl → 1.4.1__py3-none-any.whl

stoobly_agent/app/proxy/mitmproxy/response_facade.py

@@ -1,33 +1,45 @@
-from mitmproxy.http import Response as MitmproxyResponse
+import pdb
 
-from stoobly_agent.app.settings.rewrite_rule import RewriteRule 
+from mitmproxy.http import Headers, Response as MitmproxyResponse
+from mitmproxy.coretypes import multidict
+from typing import Callable, List
 
+from stoobly_agent.app.settings.constants import request_component
+from stoobly_agent.app.settings.rewrite_rule import ParameterRule, RewriteRule
+from stoobly_agent.config.constants import custom_headers
+from stoobly_agent.lib.logger import Logger, bcolors
+from stoobly_agent.lib.utils import jmespath
+from stoobly_agent.lib.utils.decode import decode
+
+from .request_facade import MitmproxyRequestFacade
+from .response_body_facade import MitmproxyResponseBodyFacade
 from .response import Response
 
+LOG_ID = 'Response'
+
 class MitmproxyResponseFacade(Response):
 
     def __init__(self, response: MitmproxyResponse):
         self.response = response
-        self.content = response.raw_content
 
-        self.rewrite_rules = []
+        self.__body = MitmproxyResponseBodyFacade(response)
+        self.__parameter_rules: List[ParameterRule] = []
 
     @property
     def code(self):
         return self.response.status_code
 
+    @property
+    def content_type(self):
+        return self.headers.get('content-type')
+
     @property
     def headers(self):
-        return self.response.headers
+        return self.__filter_custom_headers(self.response.headers)
 
     @property
     def body(self):
-        content = self.content
-
-        if not content:
-            return b''
-
-        return content
+        return self.response.raw_content or ''
 
     @property
     def http_verison(self):
@@ -40,13 +52,73 @@ class MitmproxyResponseFacade(Response):
         # Update Content-Lenght header to decoded content length
         self.response.headers['content-length'] = str(len(self.content))
 
-    
-    def with_rewrite_rules(self, rules: RewriteRule):
+    def with_parameter_rules(self, rules: List[RewriteRule], request_facade: MitmproxyRequestFacade):
         if type(rules) == list:
-            self.rewrite_rules = rules
+            self.__parameter_rules = request_facade.select_parameter_rules(rules)
+        return self 
 
-        return self
-
-    # TODO
     def rewrite(self):
-        pass
+        rewrites = self.__parameter_rules
+
+        if len(rewrites) != 0:
+            self.__rewrite_headers(rewrites)
+            self.__rewrite_content(rewrites)
+    
+    def __apply_rewrites(self, params: dict, rewrites: List[ParameterRule], handler: Callable):
+        if len(rewrites) == 0:
+            return
+
+        for rewrite in rewrites:
+            jmespath.search(rewrite.name, params, {
+                'replacements': [handler(rewrite) if handler else rewrite.value],
+            })
+
+    def __rewrite_handler(self, rewrite: ParameterRule) -> str:
+        Logger.instance(LOG_ID).info(f"{bcolors.OKCYAN}Rewriting{bcolors.ENDC} {rewrite.type.lower()} {rewrite.name} => {rewrite.value}")
+        return rewrite.value
+
+    def __rewrite_headers(self, rewrites: List[ParameterRule]):
+        self.__apply_headers(rewrites, self.__rewrite_handler)
+
+    def __rewrite_content(self, rewrites: List[ParameterRule]):
+        self.__apply_content(rewrites, self.__rewrite_handler)
+
+    def __apply_headers(self, rewrites: List[ParameterRule], handler: Callable):
+        rewrites = list(filter(lambda rewrite: rewrite.type == request_component.RESPONSE_HEADER, rewrites))
+        self.__apply_rewrites(self.response.headers, rewrites, handler)
+
+    def __apply_content(self, rewrites: List[ParameterRule], handler: Callable):
+        rewrites = list(filter(lambda rewrite: rewrite.type == request_component.RESPONSE_PARAM, rewrites))
+        if len(rewrites) == 0:
+            return
+
+        content_type = self.content_type
+        parsed_content = self.__body.get(content_type)
+
+        if not isinstance(parsed_content, dict) and not isinstance(parsed_content, multidict.MultiDictView):
+            content_type = 'application/json'
+            self.response.headers['content-type'] = content_type
+            parsed_content = {}
+
+        self.__apply_rewrites(parsed_content, rewrites, handler)
+        self.__body.set(parsed_content, content_type)
+
+    def __filter_custom_headers(self, response_headers: Headers):
+        '''
+        Remove custom headers
+        '''
+        _response_headers = Headers(**response_headers)
+
+        headers = custom_headers.__dict__
+        for key in headers:
+            if key[0:2] == '__' and key[-2:] == '__':
+                continue
+
+            name = headers[key]
+
+            if name not in response_headers:
+                continue
+
+            _response_headers.pop(name)
+
+        return _response_headers

stoobly_agent/app/proxy/record/join_request_service.py

@@ -27,6 +27,6 @@ def join_rewritten_request(flow: MitmproxyHTTPFlow, intercept_settings: Intercep
     rewrite_rules = intercept_settings.record_rewrite_rules
 
     request.with_parameter_rules(rewrite_rules).with_url_rules(rewrite_rules).rewrite()
-    response.with_rewrite_rules(rewrite_rules).rewrite()
+    response.with_parameter_rules(rewrite_rules, request).rewrite()
 
     return join_request(request, response, intercept_settings)

stoobly_agent/app/settings/constants/request_component.py

@@ -3,6 +3,7 @@ from typing import Literal
 HEADER = 'Header'
 BODY_PARAM = 'Body Param'
 QUERY_PARAM = 'Query Param'
+RESPONSE_HEADER = 'Response Header'
 RESPONSE_PARAM = 'Response Param'
 
-RequestComponent = Literal[BODY_PARAM, HEADER, QUERY_PARAM, RESPONSE_PARAM]
+RequestComponent = Literal[BODY_PARAM, HEADER, QUERY_PARAM, RESPONSE_HEADER, RESPONSE_PARAM]

stoobly_agent/config/constants/custom_headers.py

@@ -2,28 +2,28 @@ ALIAS_RESOLVE_STRATEGY = 'X-Stoobly-Alias-Resolve-Strategy'
 CONTENT_TYPE = 'X-Stoobly-Content-Type'
 CONTENT_TYPE_TEST_RESULTS = 'test/results'
 REMOTE_PROJECT_KEY = 'X-Stoobly-Endpoints-Project-Id'
-MOCK_POLICY = 'X-Mock-Policy'
+MOCK_POLICY = 'X-Stoobly-Mock-Policy'
 MOCK_REQUEST_ID = 'X-Stoobly-Request-Id'
 MOCK_REQUEST_ENDPOINT_ID = 'X-Stoobly-Request-Endpoint-Id'
 MOCK_REQUEST_KEY = 'X-Stoobly-Request-Key'
-DO_PROXY = 'X-Do-Proxy'
+DO_PROXY = 'X-Stoobly-Do-Proxy'
 LIFECYCLE_HOOKS_PATH = 'X-Stoobly-Lifecycle-Hooks-Path'
-PROJECT_KEY = 'X-Project-Key'
-PROXY_MODE = 'X-Proxy-Mode'
+PROJECT_KEY = 'X-Stoobly-Project-Key'
+PROXY_MODE = 'X-Stoobly-Proxy-Mode'
 PUBLIC_DIRECTORY_PATH = 'X-Stoobly-Public-Directory-Path'
-RECORD_POLICY = 'X-Record-Policy'
-REPORT_KEY = 'X-Report-Key'
-REQUEST_ORIGIN = 'X-Proxy-Request-Origin'
+RECORD_POLICY = 'X-Stoobly-Record-Policy'
+REPORT_KEY = 'X-Stoobly-Report-Key'
+REQUEST_ORIGIN = 'X-Stoobly-Request-Origin'
 RESPONSE_FIXTURES_PATH = 'X-Stoobly-Response-Fixtures-Path'
-RESPONSE_ID = 'X-Response-Id'
+RESPONSE_ID = 'X-Stoobly-Response-Id'
 RESPONSE_LATENCY = 'X-Stoobly-Request-Response-Latency'
-RESPONSE_PROXY_MODE = 'X-Response-Proxy-Mode'
-SCENARIO_KEY = 'X-Scenario-Key'
-SERVICE_URL = 'X-Service-Url'
-TEST_FILTER = 'X-Test-Filter'
+RESPONSE_PROXY_MODE = 'X-Stoobly-Response-Proxy-Mode'
+SCENARIO_KEY = 'X-Stoobly-Scenario-Key'
+SERVICE_URL = 'X-Stoobly-Service-Url'
+TEST_FILTER = 'X-Stoobly-Test-Filter'
 TEST_ID = 'X-Stoobly-Test-Id'
 TEST_SAVE_RESULTS = 'X-Stoobly-Test-Save-Results'
 TEST_SKIP = 'X-Stoobly-Test-Skip'
-TEST_STRATEGY = 'X-Test-Strategy'
+TEST_STRATEGY = 'X-Stoobly-Test-Strategy'
 TRACE_ID = 'X-Stoobly-Trace-Id'
 TRACE_REQUEST_ID = 'X-Stoobly-Trace-Request-Id'

stoobly_agent/config/constants/headers.py

@@ -1,9 +1,7 @@
 ACCESS_TOKEN = 'ACCESS-TOKEN'
 CLIENT = 'CLIENT'
-DO_PROXY = 'X-DO-PROXY'
 EXPIRY = 'EXPIRY'
 PROXY_HEADERS = 'X-PROXY-HEADERS'
 REQUEST_PATH = 'X-REQUEST-PATH'
-SERVICE_URL = 'X-SERVICE-URL'
 TOKEN_TYPE = 'TOKEN-TYPE'
 UID = 'UID'