stoobly-agent 1.4.2__py3-none-any.whl → 1.5.1__py3-none-any.whl

stoobly_agent/app/proxy/mock/eval_request_service.py

@@ -5,16 +5,17 @@ import re
 from mitmproxy.http import Request as MitmproxyRequest
 from requests import Response
 from typing import List, TypedDict, Union
-from stoobly_agent.config.constants import custom_headers
 
-from stoobly_agent.lib.api.param_builder import ParamBuilder
-from stoobly_agent.lib.api.interfaces.requests import RequestResponseShowQueryParams
-from stoobly_agent.lib.logger import Logger
 from stoobly_agent.app.models.request_model import RequestModel
 from stoobly_agent.app.proxy.intercept_settings import InterceptSettings
+from stoobly_agent.app.proxy.mock.custom_not_found_response_builder import CustomNotFoundResponseBuilder
 from stoobly_agent.app.settings import Settings
 from stoobly_agent.app.settings.constants import request_component
 from stoobly_agent.app.settings.match_rule import MatchRule
+from stoobly_agent.config.constants import custom_headers, query_params as request_query_params
+from stoobly_agent.lib.api.param_builder import ParamBuilder
+from stoobly_agent.lib.api.interfaces.requests import RequestResponseShowQueryParams
+from stoobly_agent.lib.logger import Logger
 
 from .hashed_request_decorator import HashedRequestDecorator
 from .search_endpoint import inject_search_endpoint
@@ -40,10 +41,6 @@ def inject_eval_request(
         request_model, intercept_settings, request, ignored_components or [], **options 
     )
 
-###
-#
-# @param settings [Settings.mode.mock | Settings.mode.record]
-#
 def eval_request(
     request_model: RequestModel,
     intercept_settings: InterceptSettings,
@@ -52,9 +49,14 @@ def eval_request(
     **options: EvalRequestOptions
 ) -> Response:
     query_params_builder = ParamBuilder({})
-    query_params_builder.with_resource_scoping(intercept_settings.project_key, intercept_settings.scenario_key)
 
-    # Tease out API returning ignored components on not found
+    try:
+        query_params_builder.with_resource_scoping(intercept_settings.project_key, intercept_settings.scenario_key)
+    except:
+        # If project_key or scenario_key are invalid, assume custom not found
+        return CustomNotFoundResponseBuilder().build()
+
+    # Tease out API returning ignored components on custom not found
     if request_model.is_local and not options.get('retry'):
         remote_project_key = intercept_settings.parsed_remote_project_key
 
@@ -63,7 +65,7 @@ def eval_request(
             remote_project_id = remote_project_key.id
             endpoint_promise = lambda: search_endpoint(remote_project_id, request.method, request.url, ignored_components=1) 
 
-            query_params_builder.with_param('endpoint_promise', endpoint_promise)
+            query_params_builder.with_param(request_query_params.ENDPOINT_PROMISE, endpoint_promise)
 
     ignored_components = __build_ignored_components(ignored_components_list or [])
     query_params_builder.with_params(__build_request_params(request, ignored_components))
@@ -142,6 +144,9 @@ def __build_optional_params(request: MitmproxyRequest, options: EvalRequestOptio
     if custom_headers.MOCK_REQUEST_ID in headers:
         optional_params['request_id'] = headers[custom_headers.MOCK_REQUEST_ID]
 
+    if custom_headers.SESSION_ID in headers:
+        optional_params[request_query_params.SESSION_ID] = headers[custom_headers.SESSION_ID]
+
     return optional_params
 
 def __filter_by_match_rules(request: MitmproxyRequest, match_rules: List[MatchRule], query_params: RequestResponseShowQueryParams):
@@ -149,9 +154,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,7 +36,16 @@ def bad_request(flow: MitmproxyHTTPFlow, message: str):
         {'Content-Type': 'text/plain'}  # (optional) headers
     )
 
-    return False
+def enable_cors(flow: MitmproxyHTTPFlow):
+    flow.response = MitmproxyResponse.make(
+        200,
+        '',
+        {
+            'Access-Control-Allow-Origin': '*',
+            'Access-Control-Allow-Methods': 'GET, OPTIONS, POST, PATCH, PUT, DELETE',
+            'Access-Control-Allow-Headers': '*'
+        }
+    )
 
 # Without deleting this header, causes parsing issues when reading response
 def disable_transfer_encoding(response: MitmproxyResponse) -> None:

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/custom_headers.py

@@ -19,6 +19,7 @@ RESPONSE_ID = 'X-Stoobly-Response-Id'
 RESPONSE_LATENCY = 'X-Stoobly-Request-Response-Latency'
 RESPONSE_PROXY_MODE = 'X-Stoobly-Response-Proxy-Mode'
 SCENARIO_KEY = 'X-Stoobly-Scenario-Key'
+SESSION_ID = 'X-Stoobly-Session-Id'
 SERVICE_URL = 'X-Stoobly-Service-Url'
 TEST_FILTER = 'X-Stoobly-Test-Filter'
 TEST_ID = 'X-Stoobly-Test-Id'

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/query_params.py

@@ -0,0 +1,2 @@
+ENDPOINT_PROMISE = 'endpoint_promise'
+SESSION_ID = 'session_id'

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