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

stoobly_agent/app/proxy/mock/eval_fixtures_service.py

@@ -1,3 +1,4 @@
+import mimetypes
 import os
 import pdb
 import re
@@ -5,6 +6,7 @@ import re
 from io import BytesIO
 from mitmproxy.http import Request as MitmproxyRequest
 from requests import Response
+from requests.structures import CaseInsensitiveDict
 from typing import Union
 
 from stoobly_agent.lib.logger import bcolors, Logger
@@ -19,45 +21,88 @@ class Options():
 
 def eval_fixtures(request: MitmproxyRequest, **options: Options) -> Union[Response, None]:
   fixture_path = None
-  headers = {}
+  headers = CaseInsensitiveDict()
+  status_code = 200
 
   response_fixtures = options.get('response_fixtures')
-  fixture = __eval_response_fixtures(request, response_fixtures)
+  fixture: dict = __eval_response_fixtures(request, response_fixtures)
 
   if not fixture:
     public_directory_path = options.get('public_directory_path')
 
-    if public_directory_path and os.path.exists(public_directory_path):
-      static_file_path = os.path.join(public_directory_path, request.path.lstrip('/'))
+    if not public_directory_path:
+      return
 
-      if os.path.exists(static_file_path):
-        fixture_path = static_file_path
+    request_path = 'index' if request.path == '/' else request.path
+    _fixture_path = os.path.join(public_directory_path, request_path.lstrip('/'))
+    if request.headers.get('accept'):
+      fixture_path = __guess_file_path(_fixture_path, request.headers['accept'])
+
+    if not fixture_path:
+      fixture_path = _fixture_path
+      if not os.path.isfile(fixture_path):
+        return
   else:
     fixture_path = fixture.get('path')
-    headers = fixture.get('headers') or {}
+    if not fixture_path or not os.path.isfile(fixture_path):
+      return
 
-  if not fixture_path:
-    return
+    _headers = fixture.get('headers')
+    headers = CaseInsensitiveDict(_headers if isinstance(_headers, dict) else {}) 
+
+    if fixture.get('status_code'):
+      status_code = fixture.get('status_code')
 
   with open(fixture_path, 'rb') as fp:
     response = Response()
 
-    response.status_code = 200
+    response.status_code = status_code
     response.raw = BytesIO(fp.read()) 
     response.headers = headers
 
+    if not response.headers.get('content-type'):
+      content_type = __guess_content_type(fixture_path)
+      if content_type:
+        response.headers['content-type'] = content_type
+
     Logger.instance(LOG_ID).debug(f"{bcolors.OKBLUE}Resolved{bcolors.ENDC} fixture {fixture_path}")
 
     return response
 
+def __guess_content_type(file_path):
+  file_extension = os.path.splitext(file_path)[1]
+  if not file_extension:
+    return
+  return mimetypes.types_map.get(file_extension)
+
+def __guess_file_path(file_path, content_type):
+  file_extension = os.path.splitext(file_path)[1]
+  if file_extension:
+    return file_path
+
+  if not content_type:
+    return
+
+  content_types = __parse_accept_header(content_type)
+
+  for content_type in content_types:
+    file_extension = mimetypes.guess_extension(content_type)
+
+    if not file_extension:
+      continue
+
+    _file_path = f"{file_path}{file_extension}"
+    if os.path.isfile(_file_path):
+      return _file_path
+
 def __eval_response_fixtures(request: MitmproxyRequest, response_fixtures: Fixtures):
-  if not response_fixtures:
+  if not isinstance(response_fixtures, dict):
     return
 
   method = request.method
   routes = response_fixtures.get(method)
 
-  if not routes:
+  if not isinstance(routes, dict):
     return
 
   for path_pattern in routes:
@@ -65,7 +110,26 @@ def __eval_response_fixtures(request: MitmproxyRequest, response_fixtures: Fixtu
       continue
       
     fixture = routes[path_pattern]
+    if not isinstance(fixture, dict):
+      continue
+
     path = fixture.get('path')
 
-    if path and os.path.exists(path):
-      return fixture
+    if path:
+      return fixture
+
+def __parse_accept_header(accept_header):
+    types = []
+    for item in accept_header.split(","):
+        parts = item.split(";")
+        content_type = parts[0].strip()
+        q_value = 1.0  # Default quality value
+        if len(parts) > 1 and parts[1].strip().startswith("q="):
+            try:
+                q_value = float(parts[1].strip()[2:])
+            except ValueError:
+                pass  # Keep default q_value if parsing fails
+        types.append((content_type, q_value))
+
+    # Sort by quality factor in descending order
+    return [content_type for content_type, _ in sorted(types, key=lambda x: x[1], reverse=True)]

stoobly_agent/app/proxy/mock/eval_request_service.py

@@ -149,9 +149,9 @@ def __filter_by_match_rules(request: MitmproxyRequest, match_rules: List[MatchRu
     method = request.method.upper()
 
     keep = {
-        request_component.BODY_PARAM: True,
+        request_component.BODY_PARAM: False,
         request_component.HEADER: False,
-        request_component.QUERY_PARAM: True,
+        request_component.QUERY_PARAM: False,
     }
 
     match_rules = list(filter(lambda rule: method in rule.methods, match_rules))

stoobly_agent/app/proxy/record/join_request_service.py

@@ -1,10 +1,14 @@
 import pdb
+
 from mitmproxy.http import HTTPFlow as MitmproxyHTTPFlow
+from typing import List
 
 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
 
@@ -18,15 +22,10 @@ def join_request(
     # Create JoinedRequest
     return JoinedRequest(proxy_request).with_response(adapted_response)
 
-def join_rewritten_request(flow: MitmproxyHTTPFlow, intercept_settings: InterceptSettings) -> JoinedRequest:
-    # Adapt flow.request
+def join_request_from_flow(
+    flow: MitmproxyHTTPFlow, intercept_settings: InterceptSettings
+) -> JoinedRequest:
     request = MitmproxyRequestFacade(flow.request)
-
-    # Adapt flow.response
     response = MitmproxyResponseFacade(flow.response)
-    rewrite_rules = intercept_settings.record_rewrite_rules
-
-    request.with_parameter_rules(rewrite_rules).with_url_rules(rewrite_rules).rewrite()
-    response.with_parameter_rules(rewrite_rules, request).rewrite()
 
     return join_request(request, response, intercept_settings)

stoobly_agent/app/proxy/record/upload_request_service.py

@@ -17,7 +17,7 @@ from stoobly_agent.lib.logger import Logger, bcolors
 from stoobly_agent.lib.orm.request import Request
 
 from ..utils.publish_change_service import publish_requests_modified
-from .join_request_service import join_rewritten_request
+from .join_request_service import join_request_from_flow
 
 AGENT_STATUSES = {
     'REQUESTS_MODIFIED': 'requests-modified'
@@ -50,7 +50,7 @@ def upload_request(
 ):
     Logger.instance(LOG_ID).info(f"{bcolors.OKCYAN}Recording{bcolors.ENDC} {flow.request.url}")
 
-    joined_request = join_rewritten_request(flow, intercept_settings)
+    joined_request = join_request_from_flow(flow, intercept_settings)
 
     project_key = intercept_settings.project_key
     scenario_key = intercept_settings.scenario_key

stoobly_agent/app/proxy/replay/replay_request_service.py

@@ -32,6 +32,7 @@ class ReplayRequestOptions(TypedDict):
   report_key: Union[str, None] 
   request_origin: Union[request_origin.CLI, None] 
   response_fixtures_path: str
+  response_mode: Union[mode.RECORD, None]
   scenario_key: Union[str, None] 
   scheme: str
   test_filter: test_filter.TestFilter
@@ -83,6 +84,9 @@ def replay(context: ReplayContext, options: ReplayRequestOptions) -> requests.Re
   if options.get('response_fixtures_path'):
     __handle_path_header(custom_headers.RESPONSE_FIXTURES_PATH, options['response_fixtures_path'], headers)
 
+  if options.get('response_mode'):
+    headers[custom_headers.RESPONSE_PROXY_MODE] = options['response_mode']
+
   if options.get('scenario_key'):
     headers[custom_headers.SCENARIO_KEY] = options['scenario_key']
 
@@ -173,10 +177,6 @@ def __handle_mode_option(_mode, request: Request, headers):
       headers[custom_headers.MOCK_REQUEST_ID] = str(request.id)
 
     headers[custom_headers.MOCK_POLICY] = mock_policy.ALL
-  elif _mode == mode.RECORD:
-    # If recording, then it's actually a replay and record
-    headers[custom_headers.PROXY_MODE] = mode.REPLAY
-    headers[custom_headers.RESPONSE_PROXY_MODE] = mode.RECORD
 
 def __create_replayed_response(request_id: int, res: requests.Response, latency: int):   
   replayed_response_model = ReplayedResponseModel(Settings.instance())

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

@@ -11,7 +11,7 @@ from stoobly_agent.lib.logger import Logger, bcolors
 from stoobly_agent.app.proxy.intercept_settings import InterceptSettings
 
 from ...intercept_settings import InterceptSettings
-from ...record.join_request_service import join_rewritten_request
+from ...record.join_request_service import join_request_from_flow
 
 LOG_ID = 'Test'
 
@@ -43,7 +43,7 @@ def upload_test(
   flow: MitmproxyHTTPFlow, 
   **kwargs: UploadTestData
 ) -> Response:
-    joined_request = join_rewritten_request(flow, intercept_settings)
+    joined_request = join_request_from_flow(flow, intercept_settings)
 
     Logger.instance(LOG_ID).info(f"{bcolors.OKCYAN}Uploading{bcolors.ENDC} test results for {joined_request.proxy_request.url()}")
 

stoobly_agent/app/proxy/utils/allowed_request_service.py

@@ -6,7 +6,7 @@ from typing import List
 
 from stoobly_agent.app.proxy.intercept_settings import InterceptSettings
 from stoobly_agent.app.settings.firewall_rule import FirewallRule
-from stoobly_agent.config.constants import mock_policy, request_origin
+from stoobly_agent.config.constants import intercept_policy, request_origin
 from stoobly_agent.lib.logger import bcolors, Logger
 
 LOG_ID = 'Firewall'
@@ -18,8 +18,8 @@ def get_active_mode_policy(request: MitmproxyRequest, intercept_settings: Interc
     if allowed_request(request, intercept_settings):
         return intercept_settings.policy
     else:
-        # If the request path does not match accepted paths, do not mock
-        return mock_policy.NONE
+        # If the request path does not match accepted paths, do not intercept
+        return intercept_policy.NONE
 
 def allowed_request(request: MitmproxyRequest, intercept_settings: InterceptSettings) -> bool:
     # If an exclude rule(s) exists, then only requests not matching these pattern(s) are allowed

stoobly_agent/app/proxy/utils/response_handler.py

@@ -36,8 +36,6 @@ def bad_request(flow: MitmproxyHTTPFlow, message: str):
         {'Content-Type': 'text/plain'}  # (optional) headers
     )
 
-    return False
-
 # Without deleting this header, causes parsing issues when reading response
 def disable_transfer_encoding(response: MitmproxyResponse) -> None:
     header_name = 'Transfer-Encoding'

stoobly_agent/app/proxy/utils/rewrite.py

@@ -0,0 +1,72 @@
+from mitmproxy.http import HTTPFlow as MitmproxyHTTPFlow
+from typing import List
+
+from stoobly_agent.app.settings.constants.request_component import REQUEST_COMPONENTS, RESPONSE_COMPONENTS
+from stoobly_agent.app.settings.rewrite_rule import RewriteRule
+
+from ..mitmproxy.request_facade import MitmproxyRequestFacade
+from ..mitmproxy.response_facade import MitmproxyResponseFacade
+
+def select_request_rewrite_rules(rewrite_rules: List[RewriteRule]):
+  rules = []
+
+  for rewrite_rule in rewrite_rules:
+    parameter_rules = list(filter(
+      lambda parameter: parameter.type in REQUEST_COMPONENTS and parameter.name, 
+      rewrite_rule.parameter_rules or []
+    ))
+
+    if len(parameter_rules) > 0:
+      rewrite_rule = RewriteRule(rewrite_rule.to_dict())
+      rewrite_rule.url_rules = rewrite_rule.url_rules
+      rewrite_rule.parameter_rules = parameter_rules
+      rules.append(rewrite_rule)
+
+  return rules
+
+def select_response_rewrite_rules(rewrite_rules: List[RewriteRule]):
+  rules = []
+
+  for rewrite_rule in rewrite_rules:
+    parameter_rules = list(filter(
+      lambda parameter: parameter.type in RESPONSE_COMPONENTS and parameter.name, 
+      rewrite_rule.parameter_rules or []
+    ))
+
+    if len(parameter_rules) > 0:
+      rewrite_rule = RewriteRule(rewrite_rule.to_dict())
+      rewrite_rule.url_rules = [] 
+      rewrite_rule.parameter_rules = parameter_rules
+      rules.append(rewrite_rule)
+
+  return rules
+
+def rewrite_request_response(flow: MitmproxyHTTPFlow, rewrite_rules: List[RewriteRule]):
+  request = rewrite_request(flow, rewrite_rules)
+  response = rewrite_response(flow, rewrite_rules, request)
+  return request, response
+
+def rewrite_request(flow: MitmproxyHTTPFlow, rewrite_rules: List[RewriteRule]):
+  request = None
+
+  # Adapt flow.request
+  request = MitmproxyRequestFacade(flow.request)
+
+  _rewrite_rules = select_request_rewrite_rules(rewrite_rules) 
+  if len(_rewrite_rules):
+    request.with_parameter_rules(_rewrite_rules).with_url_rules(_rewrite_rules).rewrite()
+
+  return request
+
+def rewrite_response(flow: MitmproxyHTTPFlow, rewrite_rules: List[RewriteRule], request = None):
+  # Adapt flow.request
+  request = request or MitmproxyRequestFacade(flow.request)
+
+  # Adapt flow.response
+  response = MitmproxyResponseFacade(flow.response)
+
+  _rewrite_rules = select_response_rewrite_rules(rewrite_rules)
+  if len(_rewrite_rules):
+    response.with_parameter_rules(_rewrite_rules, request).rewrite()
+
+  return response

stoobly_agent/app/settings/constants/request_component.py

@@ -6,4 +6,7 @@ QUERY_PARAM = 'Query Param'
 RESPONSE_HEADER = 'Response Header'
 RESPONSE_PARAM = 'Response Param'
 
-RequestComponent = Literal[BODY_PARAM, HEADER, QUERY_PARAM, RESPONSE_HEADER, RESPONSE_PARAM]
+REQUEST_COMPONENTS = [BODY_PARAM, HEADER, QUERY_PARAM]
+RESPONSE_COMPONENTS = [RESPONSE_HEADER, RESPONSE_PARAM]
+
+RequestComponent = Literal[BODY_PARAM, HEADER, QUERY_PARAM, RESPONSE_HEADER, RESPONSE_PARAM]

stoobly_agent/cli.py

@@ -159,6 +159,7 @@ def run(**kwargs):
 @click.option('--format', type=click.Choice([RAW_FORMAT]), help='Format response')
 @click.option('-H', '--header', multiple=True, help='Pass custom header(s) to server')
 @click.option('--lifecycle-hooks-path', help='Path to lifecycle hooks script.')
+@click.option('-o', '--output', help='Write to file instead of stdout')
 @ConditionalDecorator(lambda f: click.option('--project-key')(f), is_remote)
 @click.option('--public-directory-path', help='Path to public files. Used for mocking requests.')
 @click.option('--response-fixtures-path', help='Path to response fixtures yaml. Used for mocking requests.')
@@ -169,16 +170,7 @@ def mock(**kwargs):
   if kwargs.get('remote_project_key'):
     validate_project_key(kwargs['remote_project_key'])
 
-  if kwargs.get('scenario_key'):
-    validate_scenario_key(kwargs['scenario_key'])
-
-  request = __build_request_from_curl(**kwargs)
-
-  context = ReplayContext.from_python_request(request)
-  response: requests.Response = replay_request(context, {
-    **kwargs,
-    'mode': mode.MOCK,
-  })
+  response = __replay(mode.MOCK, **kwargs)
 
   if response.status_code == custom_response_codes.NOT_FOUND:
     content = response.content
@@ -186,10 +178,15 @@ def mock(**kwargs):
     sys.exit(1)
 
   if kwargs['format'] == RAW_FORMAT:
-    print_raw_response(response)
+    print_raw_response(response, kwargs['output'])
   else:
     content = response.content
-    print(decode(content), end='')
+
+    if not kwargs['output']:
+      print(decode(content), end='')
+    else:
+      with open(kwargs['output'], 'w') as fp:
+        fp.write(decode(content)) 
 
 @main.command(
   help="Record request"
@@ -197,30 +194,28 @@ def mock(**kwargs):
 @click.option('-d', '--data', default='', help='HTTP POST data')
 @click.option('--format', type=click.Choice([RAW_FORMAT]), help='Format response')
 @click.option('-H', '--header', multiple=True, help='Pass custom header(s) to server')
+@click.option('--lifecycle-hooks-path', help='Path to lifecycle hooks script.')
+@click.option('-o', '--output', help='Write to file instead of stdout')
 @ConditionalDecorator(lambda f: click.option('--project-key')(f), is_remote)
 @click.option('-X', '--request', default='GET', help='Specify request command to use')
 @click.option('--scenario-key')
 @click.argument('url')
 def record(**kwargs):
-  if kwargs.get('scenario_key'):
-    validate_scenario_key(kwargs['scenario_key'])
-
-  request = __build_request_from_curl(**kwargs)
-
-  context = ReplayContext.from_python_request(request)
-  response: requests.Response = replay_request(context, {
-    **kwargs,
-    'mode': mode.RECORD,
-  })
+  response = __replay(mode.RECORD, **kwargs) 
 
   if kwargs['format'] == RAW_FORMAT:
-    print_raw_response(response)
+    print_raw_response(response, kwargs['output'])
   else:
-    try:
-      content = response.raw.data
-      print(content.decode(json.detect_encoding(content)), end='')
-    except UnicodeDecodeError:
-      print('Warning: Binary output can mess up your terminal.')
+    content: bytes = response.raw.data
+
+    if not kwargs['output']:
+      try:
+        print(content.decode(json.detect_encoding(content)), end='')
+      except UnicodeDecodeError:
+        print('Warning: Binary output can mess up your terminal.')
+    else:
+      with open(kwargs['output'], 'w') as fp:
+        fp.write(content.decode(json.detect_encoding(content))) 
 
 def __build_request_from_curl(**kwargs):
   headers = {}
@@ -238,3 +233,15 @@ def __build_request_from_curl(**kwargs):
     method=kwargs['request'],
     url=kwargs['url']
   )
+
+def __replay(mode, **kwargs):
+  if kwargs.get('scenario_key'):
+    validate_scenario_key(kwargs['scenario_key'])
+
+  request = __build_request_from_curl(**kwargs)
+
+  context = ReplayContext.from_python_request(request)
+  return replay_request(context, {
+    **kwargs,
+    'mode': mode,
+  })

stoobly_agent/config/constants/intercept_policy.py

@@ -0,0 +1,2 @@
+ALL = 'all'
+NONE = 'none'

stoobly_agent/config/constants/mock_policy.py

@@ -1,3 +1,5 @@
-ALL = 'all'
-NONE = 'none'
+from .intercept_policy import ALL as INTERCEPT_ALL, NONE as INTERCEPT_NONE
+
+ALL = INTERCEPT_ALL
+NONE = INTERCEPT_NONE
 FOUND = 'found'

stoobly_agent/config/constants/record_policy.py

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

stoobly_agent/config/constants/replay_policy.py

@@ -1,2 +1,4 @@
-ALL = 'all'
-NONE = 'none'
+from .intercept_policy import ALL as INTERCEPT_ALL, NONE as INTERCEPT_NONE
+
+ALL = INTERCEPT_ALL
+NONE = INTERCEPT_NONE