stoobly-agent 1.2.3__py3-none-any.whl → 1.4.0__py3-none-any.whl

stoobly_agent/app/models/helpers/apply.py

@@ -61,15 +61,17 @@ class Apply():
       return
 
     last_processed_event = None
+    completed = True
 
     for event in unprocessed_events:
       if self.__logger:
-        self.__logger(f"Processing event {event.uuid}")
+        self.__logger(f"{bcolors.OKBLUE}Processing Event{bcolors.ENDC} {event.uuid}")
 
       results = event.apply(**self.__handlers())
       if results:
         status = results[1]
-        if status == 0 or status >= 500:
+        if status == 0 or status >= 400:
+          completed = False
           break
 
       last_processed_event = event
@@ -86,6 +88,8 @@ class Apply():
 
     log.lock()
 
+    return completed
+
   def request(self, uuid: str):
     result = self.__apply_put_request(uuid)
     if not result:
@@ -121,9 +125,13 @@ class Apply():
       return False
 
     if self.__logger:
-      self.__logger(f"Processing event {event.uuid}")
+      self.__logger(f"{bcolors.OKBLUE}Processing Event{bcolors.ENDC} {event.uuid}")
 
-    event.apply(**self.__handlers())
+    results = event.apply(**self.__handlers())
+    if results:
+      status = results[1]
+      if status == 0 or status >= 400:
+        return False
 
     return True
 
@@ -139,7 +147,7 @@ class Apply():
     res, status = self.request_model.destroy(uuid, force=self.__force)
 
     if status == 200:
-      self.__logger(f"{bcolors.WARNING}Deleted{bcolors.ENDC} request {uuid}") 
+      self.__logger(f"{bcolors.WARNING}Deleted Request{bcolors.ENDC} {uuid}") 
     else: 
       self.__logger(f"{bcolors.FAIL}{status}{bcolors.ENDC} {res}")
 
@@ -150,9 +158,9 @@ class Apply():
 
     raw_request = snapshot.request
     if not raw_request:
-      error = f"Snapshot for request {uuid} not found" 
-      self.__logger(f"{bcolors.FAIL}400{bcolors.ENDC} {error}")
-      return error, 400
+      error = f"snapshot for request {uuid} not found" 
+      self.__logger(f"{bcolors.WARNING}Skipping Request{bcolors.ENDC} {error}")
+      return error, 301
 
     return self.__put_request(uuid, raw_request)
 
@@ -160,7 +168,7 @@ class Apply():
     res, status = self.scenario_model.destroy(uuid, force=self.__force)
 
     if self.__logger and status == 200:
-      self.__logger(f"{bcolors.WARNING}Deleted{bcolors.ENDC} scenario {uuid}")
+      self.__logger(f"{bcolors.WARNING}Deleted Scenario{bcolors.ENDC} {uuid}")
     else: 
       self.__logger(f"{bcolors.FAIL}{status}{bcolors.ENDC} {res}")
 
@@ -171,9 +179,9 @@ class Apply():
     metadata = snapshot.metadata
 
     if not metadata:
-      error = f"Snapshot for scenario {uuid} not found"
-      self.__logger(f"{bcolors.FAIL}400{bcolors.ENDC} {error}")
-      return error, 400
+      error = f"snapshot for scenario {uuid} not found"
+      self.__logger(f"{bcolors.WARNING}Skipping Scenario{bcolors.ENDC} {error}")
+      return error, 301 
 
     res, status = self.scenario_model.show(uuid)
     if status == 404:
@@ -184,7 +192,7 @@ class Apply():
 
       if self.__logger:
         if status == 200:
-          self.__logger(f"{bcolors.OKGREEN}Created scenario{bcolors.ENDC} {res['name']}") 
+          self.__logger(f"{bcolors.OKGREEN}Created Scenario{bcolors.ENDC} {res['name']}") 
         else: 
           self.__logger(f"{bcolors.FAIL}{status}{bcolors.ENDC} {res}")
     else:
@@ -195,7 +203,7 @@ class Apply():
 
       if self.__logger:
         if status == 200:
-          self.__logger(f"{bcolors.OKBLUE}Updated{bcolors.ENDC} scenario {res['name']}") 
+          self.__logger(f"{bcolors.OKCYAN}Updated Scenario{bcolors.ENDC} {res['name']}") 
         else:
           self.__logger(f"{bcolors.FAIL}{status}{bcolors.ENDC} {res}")
 
@@ -221,6 +229,9 @@ class Apply():
       uuid = control.id
       res, status = self.__put_request(uuid, raw_request, scenario_id=scenario['id'])
 
+      if status != 200:
+        return res, status
+
       snapshot_requests[uuid] = res
 
     # Remove requests in scenario that don't exist in the snapshot
@@ -242,15 +253,21 @@ class Apply():
     res, status = self.request_model.show(uuid)
 
     if status == 404:
+      request_params = build_params(raw_request)
+
+      if not request_params:
+          self.__logger(f"{bcolors.FAIL}{status}{bcolors.ENDC} failed to join request {uuid}")
+          return res, status
+
       params = {
-        **build_params(raw_request),
+        **request_params,
         **base_params,
       }
 
       res, status = self.request_model.create(**params)
 
       if self.__logger and status == 200:
-        self.__logger(f"{bcolors.OKGREEN}Created{bcolors.ENDC} {res['list'][0]['url']}")
+        self.__logger(f"{bcolors.OKGREEN}Created Request{bcolors.ENDC} {res['list'][0]['url']}")
       else: 
         self.__logger(f"{bcolors.FAIL}{status}{bcolors.ENDC} {res}")
     elif status == 200:
@@ -263,7 +280,7 @@ class Apply():
 
       if self.__logger:
         if status == 200:
-          self.__logger(f"{bcolors.OKBLUE}Updated{bcolors.ENDC} {res['url']}")
+          self.__logger(f"{bcolors.OKCYAN}Updated Request{bcolors.ENDC} {res['url']}")
         else: 
           self.__logger(f"{bcolors.FAIL}{status}{bcolors.ENDC} {res}")
 

stoobly_agent/app/models/helpers/create_request_params_service.py

@@ -6,6 +6,9 @@ from stoobly_agent.app.models.adapters.python import PythonRequestAdapterFactory
 
 from stoobly_agent.app.proxy.record.join_request_service import InterceptSettings, join_request, MitmproxyRequestFacade, MitmproxyResponseFacade
 from stoobly_agent.app.settings import Settings
+from stoobly_agent.lib.logger import Logger
+
+LOG_ID = 'CreateRequestParamsService'
 
 class MitmproxyFlowMock():
     def __init__(self, request, response):
@@ -16,6 +19,7 @@ def build_params(raw_requests: str, payloads_delimitter = None):
   try:
     joined_request = JoinedRequestAdapter(raw_requests, payloads_delimitter).adapt()
   except Exception as e:
+    Logger.instance(LOG_ID).error(e)
     return
 
   request_adapter = RawHttpRequestAdapter(joined_request.request_string.get())

stoobly_agent/app/proxy/handle_mock_service.py

@@ -71,6 +71,7 @@ def handle_request_mock_generic(context: MockContext, **options: MockOptions):
         context.with_response(res)
 
         if handle_success:
+            # TODO: rewrite response, see #332
             res = handle_success(context) or res
     elif policy == mock_policy.FOUND:
         res = eval_request_with_retry(context, eval_request, **options) 
@@ -86,6 +87,7 @@ def handle_request_mock_generic(context: MockContext, **options: MockOptions):
                     pass
         else:
             if handle_success:
+                # TODO: rewrite response, see #332
                 res = handle_success(context) or res
     else:
         return bad_request(

stoobly_agent/app/proxy/handle_replay_service.py

@@ -24,6 +24,8 @@ def handle_request_replay(replay_context: ReplayContext):
 def handle_response_replay(replay_context: ReplayContext):
     __replay_hook(lifecycle_hooks.AFTER_REPLAY, replay_context)
 
+    # TODO: rewrite response, see #332
+
 def __replay_request(replay_context: ReplayContext):
     """
     Before replaying a request, see if the request needs to be rewritten

stoobly_agent/app/proxy/mitmproxy/request_facade.py

@@ -70,7 +70,7 @@ class MitmproxyRequestFacade(Request):
 
     @property
     def body(self):
-        content = self.request.raw_content or ''
+        content = self.request.raw_content or b''
 
         return decode(content)
 

stoobly_agent/app/proxy/mitmproxy/response_body_facade.py

@@ -0,0 +1,19 @@
+import pdb
+
+from mitmproxy.http import Response as MitmproxyResponse
+from typing import Union
+
+from ..replay.body_parser_service import decode_response, encode_response
+
+class MitmproxyResponseBodyFacade:
+    def __init__(self, response: MitmproxyResponse):
+        self.__response = response
+
+    def get(self, content_type: Union[bytes, str]):
+        return decode_response(self.__response.content, content_type)
+
+    def set(self, content, content_type: Union[bytes, str]):
+        """
+        Adjusting Content-Length header should be done by MitmproxyResponse
+        """
+        self.__response.content = encode_response(content, content_type).encode()

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/proxy/replay/body_parser_service.py

@@ -1,8 +1,8 @@
-import cgi
 import json
 import pdb
 import urllib.parse
 
+from email.message import Message
 from mitmproxy.coretypes.multidict import MultiDict
 from mitmproxy.net import encoding
 from typing import Dict, Union
@@ -108,11 +108,19 @@ def serialize_www_form_urlencoded(o):
 def normalize_header(header):
     if isinstance(header, bytes):
         header = header.decode('utf-8')
-    return cgi.parse_header(header)[0].lower()
+    return __parse_separated_header(header).lower()
 
 def is_traversable(content):
     return isinstance(content, list) or isinstance(content, dict) or isinstance(content, MultiDict)
 
 def is_json(content_type):
    _content_type = content_type.lower()
-   return _content_type == JSON or _content_type.startswith('application/x-amz-json') 
+   return _content_type == JSON or _content_type.startswith('application/x-amz-json') 
+
+
+def __parse_separated_header(header: str):
+    # Adapted from https://peps.python.org/pep-0594/#cgi
+    message = Message()
+    message['content-type'] = header 
+    return message.get_content_type()
+

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'

stoobly_agent/config/data_dir.py

@@ -4,6 +4,7 @@ import shutil
 
 from stoobly_agent.config.constants.env_vars import ENV
 
+CERTS_DIR_NAME = 'certs'
 DATA_DIR_NAME = '.stoobly'
 DB_FILE_NAME = 'stoobly_agent.sqlite3'
 DB_VERSION_NAME = 'VERSION'
@@ -85,7 +86,7 @@ class DataDir:
 
     @property
     def certs_dir_path(self):
-        certs_dir_path = os.path.join(self.path, 'certs')
+        certs_dir_path = os.path.join(self.path, CERTS_DIR_NAME)
 
         if not os.path.exists(certs_dir_path):
             os.mkdir(certs_dir_path)

stoobly_agent/config/schema.yml

@@ -6,13 +6,13 @@ cli:
 proxy:
   data: map(include('data_rules'), key=include('project_id'))
   firewall: map(list(include('firewall_rule')), key=include('project_id'))
-  match: map(list(include('match_rule')), key=include('project_id'))
+  match: map(list(include('match_rule')), key=include('project_id'), required=False)
   rewrite: map(list(include('rewrite_rule')), key=include('project_id'))
   intercept:
     active: bool()
     mode: include('mode')
     project_key: str()
-  url: str()
+  url: str(required=False)
 proxy_config_path: str(required=False)
 remote: 
   api_url: str()