stoobly-agent 1.8.4__py3-none-any.whl → 1.9.0__py3-none-any.whl

stoobly_agent/app/proxy/handle_mock_service.py

@@ -3,7 +3,7 @@ import pdb
 import requests
 import time
 
-from mitmproxy.http import HTTPFlow as MitmproxyHTTPFlow, Request as MitmproxyRequest
+from mitmproxy.http import Request as MitmproxyRequest
 from typing import Callable, TypedDict
 
 from stoobly_agent.app.models.request_model import RequestModel
@@ -64,7 +64,7 @@ def handle_request_mock_generic(context: MockContext, **options: MockOptions):
 
         return bad_request(
             context.flow,
-            "Valid env MOCK_POLICY: %s, Got: %s" %
+            "Valid mock policies: %s, Got: %s" %
             ([mock_policy.ALL, mock_policy.FOUND, mock_policy.NONE], policy)
         )
 
@@ -88,18 +88,20 @@ def handle_request_mock_generic(context: MockContext, **options: MockOptions):
         res = eval_request_with_retry(context, eval_request, **options) 
 
         context.with_response(res)
+        res = __after_mock_not_found(context)
     elif policy == mock_policy.FOUND:
         res = eval_request_with_retry(context, eval_request, **options) 
 
         context.with_response(res)
+        res = __after_mock_not_found(context)
 
-        if res.status_code in [custom_response_codes.NOT_FOUND, custom_response_codes.IGNORE_COMPONENTS]:
+        if res.status_code == custom_response_codes.NOT_FOUND:
             try:
                 return __handle_found_policy(context) # Continue proxying the request
             except RuntimeError:
                 # Do nothing, return custom error response
-                pass
-    
+                pass 
+
     if res.status_code == custom_response_codes.NOT_FOUND:
         if handle_failure:
             res = handle_failure(context) or res
@@ -241,4 +243,15 @@ def __mock_hook(hook: str, context: MockContext):
     lifecycle_hooks_module = intercept_settings.lifecycle_hooks
 
     if hook in lifecycle_hooks_module:
-        lifecycle_hooks_module[hook](context)
+        lifecycle_hooks_module[hook](context)
+
+def __after_mock_not_found(context: MockContext):
+    res = context.response
+
+    if res.status_code == custom_response_codes.NOT_FOUND:
+        __mock_hook(lifecycle_hooks.AFTER_MOCK_NOT_FOUND, context)
+
+        # context.response may have been modified by the hook
+        res = context.response
+
+    return res

stoobly_agent/app/proxy/handle_record_service.py

@@ -9,7 +9,7 @@ from stoobly_agent.app.settings.constants.mode import TEST
 from stoobly_agent.app.models.request_model import RequestModel
 from stoobly_agent.app.proxy.intercept_settings import InterceptSettings
 from stoobly_agent.config.constants.env_vars import ENV
-from stoobly_agent.config.constants import lifecycle_hooks, record_policy
+from stoobly_agent.config.constants import lifecycle_hooks, record_order, record_policy
 from stoobly_agent.lib.logger import Logger
 
 from .constants import custom_response_codes
@@ -17,6 +17,7 @@ from .mock.eval_request_service import inject_eval_request
 from .record.context import RecordContext
 from .record.overwrite_scenario_service import overwrite_scenario
 from .record.upload_request_service import inject_upload_request
+from .replay.body_parser_service import is_json, is_xml
 from .utils.allowed_request_service import get_active_mode_policy
 from .utils.response_handler import bad_request, disable_transfer_encoding 
 from .utils.rewrite import rewrite_request_response
@@ -42,6 +43,13 @@ def handle_response_record(context: RecordContext):
 
     if active_record_policy == record_policy.ALL:
         __record_request(context, request_model)
+    elif active_record_policy == record_policy.API:
+        response = flow.response
+        content_type: str = response.headers.get('content-type')
+
+        if content_type:
+            if is_json(content_type) or is_xml(content_type) or content_type.startswith('text/plain'):
+                __record_request(context, request_model)
     elif active_record_policy == record_policy.FOUND:
         res = inject_eval_request(request_model, intercept_settings)(request, [])
 
@@ -52,16 +60,12 @@ def handle_response_record(context: RecordContext):
 
         if res.status_code == custom_response_codes.NOT_FOUND:
             __record_request(context, request_model)
-    elif active_record_policy == record_policy.OVERWRITE:
-        overwrite_scenario(intercept_settings.scenario_key)
-
-        __record_request(context, request_model)
     else:
         if active_record_policy != record_policy.NONE:
             return bad_request(
                 flow,
-                "Valid env RECORD_POLICY: %s, Got: %s" %
-                ([record_policy.ALL, record_policy.FOUND, record_policy.NOT_FOUND], active_record_policy)
+                "Valid record policies: %s, Got: %s" %
+                ([record_policy.ALL, record_policy.API, record_policy.FOUND, record_policy.NOT_FOUND], active_record_policy)
             )
 
 def __record_handler(context: RecordContext, request_model: RequestModel):
@@ -79,6 +83,11 @@ def __record_handler(context: RecordContext, request_model: RequestModel):
     context.flow = flow # Reset flow
 
 def __record_request(context: RecordContext, request_model: RequestModel):
+    intercept_settings = context.intercept_settings
+
+    if intercept_settings.order == record_order.OVERWRITE:
+        overwrite_scenario(intercept_settings.scenario_key)
+
     if os.environ.get(ENV) == TEST:
         __record_handler(context, request_model)
     else:

stoobly_agent/app/proxy/intercept_settings.py

@@ -166,6 +166,10 @@ class InterceptSettings:
     if self.__headers and custom_headers.REPORT_KEY in self.__headers:
       return self.__headers[custom_headers.REPORT_KEY]
 
+  @property
+  def order(self):
+    return self.__order(self.mode)
+
   @property
   def policy(self):
     return self.__policy(self.mode)
@@ -339,7 +343,11 @@ class InterceptSettings:
             self.__response_fixtures = yaml.safe_load(stream)
         except yaml.YAMLError as exc:
             Logger.instance().error(exc)
-        
+
+  def __order(self, mode):
+    if mode == intercept_mode.RECORD:
+      return self.__data_rules.record_order
+    
   def __policy(self, mode):
     if mode == intercept_mode.MOCK:
       if self.__headers and custom_headers.MOCK_POLICY in self.__headers:
@@ -352,6 +360,9 @@ class InterceptSettings:
 
       return self.__data_rules.record_policy
     elif mode == intercept_mode.TEST:
+      if self.__headers and custom_headers.TEST_POLICY in self.__headers:
+        return self.__headers[custom_headers.TEST_POLICY]
+
       return self.__data_rules.test_policy
     elif mode == intercept_mode.REPLAY:
       return self.__data_rules.replay_policy

stoobly_agent/app/proxy/mitmproxy/request_facade.py

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

stoobly_agent/app/proxy/record/join_request_service.py

@@ -1,31 +1,43 @@
 import pdb
 
-from mitmproxy.http import HTTPFlow as MitmproxyHTTPFlow
-from typing import List
+from mitmproxy.http import HTTPFlow as MitmproxyHTTPFlow, Request as MitmproxyRequest, Response as MitmproxyResponse
+from typing import TypedDict
 
 from stoobly_agent.app.proxy.intercept_settings import InterceptSettings
-from stoobly_agent.app.settings.rewrite_rule import RewriteRule
 
 from ..mitmproxy.request_facade import MitmproxyRequestFacade
 from ..mitmproxy.response_facade import MitmproxyResponseFacade
-from ..utils.rewrite import rewrite_request_response
 from .joined_request import JoinedRequest
 from .proxy_request import ProxyRequest
 
+class JoinRequestOptions(TypedDict):
+    id: str
+    intercept_settings: InterceptSettings
+
 def join_request(
-    adapted_request: MitmproxyRequestFacade, adapted_response: MitmproxyResponseFacade, intercept_settings: InterceptSettings
+    adapted_request: MitmproxyRequestFacade, adapted_response: MitmproxyResponseFacade, **options: JoinRequestOptions
 ) -> JoinedRequest:
+    intercept_settings: InterceptSettings = options.get('intercept_settings')
+
     # Decorate request with service_url
-    upstream_url = intercept_settings.upstream_url
+    upstream_url = intercept_settings.upstream_url if intercept_settings else None
     proxy_request = ProxyRequest(adapted_request, upstream_url)
 
+    if options.get('id'):
+        proxy_request.id = options['id']
+
     # Create JoinedRequest
     return JoinedRequest(proxy_request).with_response(adapted_response)
 
 def join_request_from_flow(
-    flow: MitmproxyHTTPFlow, intercept_settings: InterceptSettings
+    flow: MitmproxyHTTPFlow, **options: JoinRequestOptions
 ) -> JoinedRequest:
-    request = MitmproxyRequestFacade(flow.request)
-    response = MitmproxyResponseFacade(flow.response)
+    return join_request_from_request_response(flow.request, flow.response, **options)
+
+def join_request_from_request_response(
+    request: MitmproxyRequest, response: MitmproxyResponse, **options: JoinRequestOptions
+):
+    request = MitmproxyRequestFacade(request)
+    response = MitmproxyResponseFacade(response)
 
-    return join_request(request, response, intercept_settings)
+    return join_request(request, response, **options)

stoobly_agent/app/proxy/record/proxy_request.py

@@ -2,10 +2,19 @@ from ..mitmproxy.request_facade import MitmproxyRequestFacade
 
 class ProxyRequest:
     def __init__(self, request: MitmproxyRequestFacade, upstream_url: str = None):
+        self._id = None
         self.request = request
 
         self.upstream_url = upstream_url
 
+    @property
+    def id(self):
+        return self._id
+
+    @id.setter
+    def id(self, v):
+        self._id = v
+
     def url(self):
         url = self.request.url
 

stoobly_agent/app/proxy/record/request_string.py

@@ -9,12 +9,13 @@ from stoobly_agent.lib.utils.decode import decode
 from .proxy_request import ProxyRequest
 from .request_string_control import RequestStringControl
 
-CLRF = "\r\n"
+CLRF = b"\r\n"
 
 class RequestString:
     ENCODING = 'utf-8'
 
     __current_time = None
+    request_id = None
 
     def __init__(self, proxy_request: ProxyRequest):
         self.__current_time = self.__get_current_time()
@@ -29,24 +30,25 @@ class RequestString:
             self.__headers()
             self.__body()
 
-        self.request_id = self.__generate_request_id()
+            self.request_id = proxy_request.id
+
+        self.request_id = self.request_id or self.__generate_request_id()
 
     def get(self, **kwargs):
         if kwargs.get('control'):
-            return CLRF.join([self.control] + self.lines).encode(self.ENCODING)
+            return CLRF.join([self.control] + self.lines)
         else:
-            return CLRF.join(self.lines).encode(self.ENCODING)
+            return CLRF.join(self.lines)
 
     def set(self, s: bytes):
-        decoded_s = decode(s, self.ENCODING)
-        self.lines = decoded_s.split(CLRF)
+        self.lines = s.split(CLRF)
 
     @property
     def control(self):
         control = RequestStringControl()
         control.id = self.request_id
         control.timestamp = self.__current_time
-        return control.serialize()
+        return control.serialize().encode(self.ENCODING)
 
     @control.setter
     def control(self, c: str):
@@ -55,20 +57,27 @@ class RequestString:
         self.__current_time = control.timestamp
 
     def __request_line(self):
-        self.lines.append("{} {} HTTP/1.1".format(self.request.method, self.proxy_request.url()))
+        line = "{} {} HTTP/1.1".format(self.request.method, self.proxy_request.url())
+        self.lines.append(line.encode(self.ENCODING))
 
     def __headers(self):
         headers = self.request.headers
 
         for name, val in headers.items():
+            if 'content-encoding' in headers:
+                # self.request.body will decode content-encoding
+                # to maintain internal consistency, remove this header since it will be decoded
+                print(self.request.body)
+                pass
+
             line = ' '.join([
                 "{}:".format(self.__to_header_case(self.__to_str(name))), 
                 self.__to_str(val)
             ])
-            self.lines.append(line)
+            self.lines.append(line.encode(self.ENCODING))
 
     def __body(self):
-        self.lines.append("{}{}".format(CLRF, self.request.body))
+        self.lines.append(CLRF + self.request.body)
 
     def __to_header_case(self, header: str) -> str:
         toks = header.split('_')

stoobly_agent/app/proxy/record/upload_request_service.py

@@ -4,6 +4,7 @@ import pdb
 import time
 import tempfile
 
+from copy import deepcopy
 from mitmproxy.http import HTTPFlow as MitmproxyHTTPFlow
 from mitmproxy.http import Request as MitmproxyRequest
 
@@ -50,22 +51,23 @@ def upload_request(
 ):
     Logger.instance(LOG_ID).info(f"{bcolors.OKCYAN}Recording{bcolors.ENDC} {flow.request.url}")
 
-    joined_request = join_request_from_flow(flow, intercept_settings)
+    flow_copy = deepcopy(flow) # When applying modifications we don't want to persist them in the response
+    joined_request = join_request_from_flow(flow_copy, intercept_settings=intercept_settings)
 
     project_key = intercept_settings.project_key
     scenario_key = intercept_settings.scenario_key
     body_params = __build_body_params(
         project_key,
-        joined_request, 
+        joined_request,
         flow=flow,
         scenario_key=scenario_key
     )
 
     # If request_origin is WEB, then we are in proxy
     # This means that we have access to Cache singleton and do not need send a request to update the status
-    sync = intercept_settings.request_origin == request_origin.WEB 
+    sync = intercept_settings.request_origin == request_origin.WEB
     res = __upload_request_with_body_params(request_model, body_params, sync)
- 
+
     if intercept_settings.settings.is_debug():
         file_path = __debug_request(flow.request, joined_request.build())
         Logger.instance(LOG_ID).debug(f"Writing request to {file_path}")
@@ -142,4 +144,4 @@ def __debug_request(request: MitmproxyRequest, raw_requests: bytes):
     with open(file_path, 'wb') as f:
         f.write(raw_requests)
 
-    return file_path
+    return file_path

stoobly_agent/app/proxy/replay/body_parser_service.py

@@ -14,6 +14,7 @@ from .multipart import decode as multipart_decode, encode as multipart_encode
 JSON = 'application/json'
 MULTIPART_FORM = 'multipart/form-data'
 WWW_FORM_URLENCODED = 'application/x-www-form-urlencoded'
+XML = 'application/xml'
 
 def compress(body: Union[bytes, str], content_encoding: Union[None, str]) -> Union[bytes, str]:
     if content_encoding:
@@ -113,14 +114,24 @@ def normalize_header(header):
 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') 
+def is_json(content_type: str):
+    if not content_type:
+        return False
+
+    _content_type = content_type.lower().split(';')[0]
+    # e.g. custom json content-type: application/x-amz-json
+    return _content_type == JSON or (_content_type.startswith('application/') and _content_type.endswith('json'))
 
+def is_xml(content_type: str):
+    if not content_type:
+        return False
+
+    _content_type = content_type.lower().split(';')[0]
+    # e.g. custom json content-type: application/x-amz-json
+    return _content_type == XML or (_content_type.startswith('application/') and _content_type.endswith('xml'))
 
 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/proxy/replay/replay_request_service.py

@@ -150,6 +150,7 @@ def replay(context: ReplayContext, options: ReplayRequestOptions) -> requests.Re
     if options.get('save') or options.get('overwrite'):
       replayed_response, status = __create_replayed_response(context.request.id, res, latency)
 
+      # TODO: allow overwriting if status is the same
       if status < 400 and options.get('overwrite'):
         __overwrite_response(replayed_response.get('id'))
 

stoobly_agent/app/proxy/test/helpers/upload_test_service.py

@@ -43,7 +43,7 @@ def upload_test(
   flow: MitmproxyHTTPFlow, 
   **kwargs: UploadTestData
 ) -> Response:
-    joined_request = join_request_from_flow(flow, intercept_settings)
+    joined_request = join_request_from_flow(flow, intercept_settings=intercept_settings)
 
     Logger.instance(LOG_ID).info(f"{bcolors.OKCYAN}Uploading{bcolors.ENDC} test results for {joined_request.proxy_request.url()}")
 

stoobly_agent/app/settings/data_rules.py

@@ -1,4 +1,4 @@
-from stoobly_agent.config.constants import mock_policy, record_policy, replay_policy, test_strategy
+from stoobly_agent.config.constants import mock_policy, record_order, record_policy, replay_policy, test_strategy
 
 from .types.proxy_settings import DataRules as IDataRules
 
@@ -8,6 +8,7 @@ class DataRules:
     self.__data_rules = data_rules or {}
 
     self.__mock_policy = self.__data_rules.get('mock_policy') or mock_policy.FOUND
+    self.__record_order = self.__data_rules.get('record_order') or record_order.APPEND
     self.__record_policy = self.__data_rules.get('record_policy') or record_policy.ALL
     self.__replay_policy = self.__data_rules.get('replay_policy') or replay_policy.ALL
     self.__scenario_key = self.__data_rules.get('scenario_key')
@@ -32,6 +33,19 @@ class DataRules:
     self.__record_policy = v
     self.__data_rules['record_policy'] = v
 
+  @property
+  def record_order(self):
+    return self.__record_order
+
+  @record_order.setter
+  def record_order(self, v):
+    valid_orders = [record_order.APPEND, record_order.OVERWRITE]
+    if v not in valid_orders:
+      raise TypeError(f"record_order has to be one of {valid_orders}, got {v}")
+
+    self.__record_order = v
+    self.__data_rules['record_order'] = v
+
   @property
   def replay_policy(self):
     return self.__replay_policy
@@ -68,6 +82,7 @@ class DataRules:
   def to_dict(self) -> IDataRules:
     return {
       'mock_policy': self.__mock_policy,
+      'record_order': self.__record_order,
       'record_policy': self.__record_policy,
       'replay_policy': self.__replay_policy,
       'scenario_key': self.__scenario_key,

stoobly_agent/config/constants/custom_headers.py

@@ -23,6 +23,7 @@ SESSION_ID = 'X-Stoobly-Session-Id'
 SERVICE_URL = 'X-Stoobly-Service-Url'
 TEST_FILTER = 'X-Stoobly-Test-Filter'
 TEST_ID = 'X-Stoobly-Test-Id'
+TEST_POLICY = 'X-Stoobly-Test-Policy'
 TEST_SAVE_RESULTS = 'X-Stoobly-Test-Save-Results'
 TEST_SKIP = 'X-Stoobly-Test-Skip'
 TEST_STRATEGY = 'X-Stoobly-Test-Strategy'

stoobly_agent/config/constants/lifecycle_hooks.py

@@ -1,8 +1,10 @@
 AFTER_MOCK = 'handle_after_mock'
+AFTER_MOCK_NOT_FOUND = 'handle_after_mock_not_found'
 AFTER_RECORD = 'handle_after_record'
 AFTER_REPLAY = 'handle_after_replay'
 AFTER_TEST = 'handle_after_test'
 BEFORE_MOCK = 'handle_before_mock'
+BEFORE_MIGRATE = 'handle_before_migrate'
 BEFORE_RECORD = 'handle_before_record'
 BEFORE_REPLAY = 'handle_before_replay'
 BEFORE_REQUEST = 'handle_before_request'

stoobly_agent/config/constants/record_order.py

@@ -0,0 +1,6 @@
+from typing import Literal
+
+APPEND = 'append'
+OVERWRITE = 'overwrite'
+
+RecordOrder = Literal[APPEND, OVERWRITE]

stoobly_agent/config/constants/record_policy.py

@@ -1,7 +1,7 @@
 from .intercept_policy import ALL as INTERCEPT_ALL, NONE as INTERCEPT_NONE
 
 ALL = INTERCEPT_ALL
+API = 'api'
 FOUND = 'found'
 NONE = INTERCEPT_NONE
-NOT_FOUND = 'not_found'
-OVERWRITE = 'overwrite'
+NOT_FOUND = 'not_found'

stoobly_agent/lib/orm/transformers/orm_to_request_transformer.py

@@ -94,7 +94,7 @@ class ORMToRequestTransformer():
     return self
 
   def with_scheme(self, scheme: str):
-    if scheme != 'https' or 'http':
+    if scheme != 'https' and scheme != 'http':
       raise ValueError('Scheme must be https or http')
 
     self.__scheme = scheme