stoobly-agent 0.29.1__tar.gz → 0.30.0__tar.gz

{stoobly-agent-0.29.1/stoobly_agent.egg-info → stoobly-agent-0.30.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: stoobly-agent
-Version: 0.29.1
+Version: 0.30.0
 Summary: Client agent for Stoobly
 Home-page: https://github.com/Stoobly/stoobly-agent
 Author: Michael Yen

{stoobly-agent-0.29.1 → stoobly-agent-0.30.0}/stoobly_agent/__init__.py

@@ -1,2 +1,2 @@
 COMMAND = 'stoobly-agent'
-VERSION = '0.29.1'
+VERSION = '0.30.0'

{stoobly-agent-0.29.1 → stoobly-agent-0.30.0}/stoobly_agent/app/api/configs_controller.py

@@ -3,7 +3,9 @@ import pdb
 from mergedeep import merge
 
 from stoobly_agent.app.api.simple_http_request_handler import SimpleHTTPRequestHandler
-from stoobly_agent.app.cli.helpers.handle_config_update_service import context as handle_context, handle_project_update, handle_policy_update, handle_scenario_update
+from stoobly_agent.app.cli.helpers.handle_config_update_service import (
+    context as handle_context, handle_intercept_active_update, handle_policy_update, handle_project_update, handle_scenario_update
+) 
 from stoobly_agent.app.models.scenario_model import ScenarioModel
 from stoobly_agent.app.proxy.intercept_settings import InterceptSettings
 from stoobly_agent.app.settings import Settings
@@ -106,8 +108,9 @@ class ConfigsController:
 
         _handle_context = handle_context()
 
-        handle_project_update(settings, _handle_context)
+        handle_intercept_active_update(settings, _handle_context)
         handle_policy_update(settings, _handle_context)
+        handle_project_update(settings, _handle_context)
         handle_scenario_update(settings, _handle_context)
 
         settings.write(settings.to_dict())

{stoobly-agent-0.29.1 → stoobly-agent-0.30.0}/stoobly_agent/app/cli/config_cli.py

@@ -8,7 +8,7 @@ from stoobly_agent.app.settings import Settings
 from stoobly_agent.app.settings.constants import firewall_action, request_component
 from stoobly_agent.app.settings.firewall_rule import FirewallRule
 from stoobly_agent.app.settings.match_rule import MatchRule
-from stoobly_agent.app.settings.rewrite_rule import ParameterRule, RewriteRule
+from stoobly_agent.app.settings.rewrite_rule import ParameterRule, RewriteRule, UrlRule
 from stoobly_agent.config.constants import mode
 from stoobly_agent.config.data_dir import DataDir
 from stoobly_agent.lib.api.keys import ProjectKey, ScenarioKey
@@ -146,6 +146,7 @@ def rewrite(ctx):
 @rewrite.command(
     help="Set rewrite rule."
 )
+@click.option('--host', help='Request URL host.')
 @click.option(
     '--method', 
     multiple=True, 
@@ -159,17 +160,30 @@ def rewrite(ctx):
     required=True,
     type=click.Choice([mode.MOCK, mode.RECORD, mode.REPLAY] + ([mode.TEST] if is_remote else []))
 )
-@click.option('--name', required=True, help='Name of the request component.')
+@click.option('--name', help='Name of the request component.')
 @click.option('--pattern', required=True, help='URLs pattern.')
+@click.option('--port', help='Request URL port.')
 @click.option('--project-key', help='Project to add rewrite rule to.')
 @click.option(
     '--type', 
-    required=True,
     type=click.Choice([request_component.BODY_PARAM, request_component.HEADER, request_component.QUERY_PARAM]), 
     help='Request component type.'
 )
-@click.option('--value', required=True, help='Rewrite value.')
+@click.option('--value', help='Rewrite value.')
 def set(**kwargs):
+    if kwargs['name'] or kwargs['value'] or kwargs['type']:
+        if kwargs['name'] == None:
+            print("Error: missing option '--name'", file=sys.stderr)
+            sys.exit(1)
+
+        if kwargs['value'] == None:
+            print("Error: missing option '--value'", file=sys.stderr)
+            sys.exit(1)
+        
+        if kwargs['type'] == None:
+            print("Error: missing option '--type'", file=sys.stderr)
+            sys.exit(1)
+
     settings = Settings.instance()
     project_key_str = resolve_project_key_and_validate(kwargs, settings)
     project_key = ProjectKey(project_key_str)
@@ -183,27 +197,50 @@ def set(**kwargs):
     filtered_rewrite_rules: List[RewriteRule] = list(filter(rewrite_rule_filter, rewrite_rules))
 
     if len(filtered_rewrite_rules) == 0:
+        parameter_rules = list(filter(lambda r: r != None, [__select_parameter_rule(kwargs)]))
+        url_rules = list(filter(lambda r: r != None , [__select_url_rule(kwargs)]))
         rewrite_rule = RewriteRule({
             'methods': methods,
             'pattern': kwargs['pattern'],
-            'parameter_rules': [__select_parameter_rule(kwargs)]
+            'parameter_rules': parameter_rules,
+            'url_rules': url_rules,
         })
         rewrite_rules.append(rewrite_rule)
         settings.proxy.rewrite.set_rewrite_rules(project_key.id, rewrite_rules)
     else:
         parameter_rule_filter = lambda rule: rule.name == kwargs['name'] and rule.type == kwargs['type'] and rule.modes == modes
+        url_rule_filter = lambda rule: rule.host == kwargs['host'] and rule.modes == modes
+
         for rewrite_rule in filtered_rewrite_rules:
-            parameter_rules = rewrite_rule.parameter_rules
-            filtered_parameter_rules: List[ParameterRule] = list(filter(parameter_rule_filter, parameter_rules))
+            # Parameter rules
             parameter_rule_dict = __select_parameter_rule(kwargs)
 
-            if len(filtered_parameter_rules) == 0:
-                parameter_rule = ParameterRule(parameter_rule_dict)
-                parameter_rules.append(parameter_rule)
-                rewrite_rule.parameter_rules = parameter_rules
-            else:
-                for parameter_rule in filtered_parameter_rules:
-                    parameter_rule.update(parameter_rule_dict)
+            if parameter_rule_dict:
+                parameter_rules = rewrite_rule.parameter_rules
+                filtered_parameter_rules: List[ParameterRule] = list(filter(parameter_rule_filter, parameter_rules))
+
+                if len(filtered_parameter_rules) == 0:
+                    parameter_rule = ParameterRule(parameter_rule_dict)
+                    parameter_rules.append(parameter_rule)
+                    rewrite_rule.parameter_rules = parameter_rules
+                else:
+                    for parameter_rule in filtered_parameter_rules:
+                        parameter_rule.update(parameter_rule_dict)
+
+            # URL rules
+            url_rule_dict = __select_url_rule(kwargs)
+            
+            if url_rule_dict:
+                url_rules = rewrite_rule.url_rules
+                filtered_url_rules: List[UrlRule] = list(filter(url_rule_filter, url_rules))
+
+                if len(filtered_url_rules) == 0:
+                    url_rule = UrlRule(url_rule_dict)
+                    url_rules.append(url_rule)
+                    rewrite_rule.url_rules = url_rules
+                else:
+                    for url_rule in filtered_url_rules:
+                        url_rule.update(url_rule_dict)
 
     settings.commit()
 
@@ -437,6 +474,9 @@ config.add_command(rewrite)
 config.add_command(scenario)
 
 def __select_parameter_rule(kwargs):
+    if kwargs['name'] == None or kwargs['value'] == None or kwargs['type'] == None:
+        return
+
     return {
         'modes': list(kwargs['mode']),
         'name': kwargs['name'],
@@ -444,6 +484,16 @@ def __select_parameter_rule(kwargs):
         'type': kwargs['type'],
     }
 
+def __select_url_rule(kwargs):
+    if kwargs['host'] == None or kwargs['port'] == None:
+        return
+
+    return {
+        'host': kwargs['host'],
+        'modes': list(kwargs['mode']),
+        'port': kwargs['port'],
+    }
+
 def __project_key(settings):
     project_key = settings.proxy.intercept.project_key
     validate_project_key(project_key)

{stoobly-agent-0.29.1 → stoobly-agent-0.30.0}/stoobly_agent/app/cli/endpoint_cli.py

@@ -26,7 +26,7 @@ def endpoint(ctx):
 )
 @click.option('--format', required=True, type=click.Choice([OPENAPI_FORMAT]), help='File format')
 @click.option('--scenario-key', help='Which scenario to import to. If none then all requests will be imported to.')
-@click.option('--lifecycle-hooks-script-path', help='Path to lifecycle hooks script.')
+@click.option('--lifecycle-hooks-path', help='Path to lifecycle hooks script.')
 @ConditionalDecorator(lambda f: click.option('--project-key', help='Project to create endpoint in.')(f), is_remote)
 @click.argument('path')
 def _import(**kwargs: EndpointCreateCliOptions):

{stoobly-agent-0.29.1 → stoobly-agent-0.30.0}/stoobly_agent/app/cli/helpers/endpoint_facade.py

@@ -23,7 +23,7 @@ class EndpointFacade():
       lifecycle_hooks = {}
 
       try:
-        lifecycle_hooks = run_path(kwargs['lifecycle_hooks_script_path'])
+        lifecycle_hooks = run_path(kwargs['lifecycle_hooks_path'])
       except Exception as e:
         pass
 

{stoobly-agent-0.29.1 → stoobly-agent-0.30.0}/stoobly_agent/app/cli/helpers/handle_config_update_service.py

@@ -19,6 +19,44 @@ def context() -> Context:
     'current_proxy_settings': __current_proxy_settings()
   }
 
+def handle_intercept_active_update(new_settings: Settings, context: Context = None):
+  data_rule = __data_rule(new_settings.proxy)
+  is_active = new_settings.proxy.intercept.active
+  old_proxy_settings = __current_proxy_settings(context)
+  was_active = old_proxy_settings.intercept.active
+  _mode = context['mode'] if context and context.get('mode') else new_settings.proxy.intercept.mode
+
+  if not was_active and is_active:
+    if _mode == intercept_mode.RECORD:
+      new_policy = data_rule.record_policy
+      scenario_key = data_rule.scenario_key
+      _scenario_key = __parse_scenario_key(scenario_key)
+
+      if _scenario_key:
+        scenario_model = ScenarioModel(new_settings)
+
+        if new_policy == record_policy.OVERWRITE:
+          # If policy is overwrite when recording, whenever intercept is enabled,
+          # set active scenario to be overwritable
+          scenario_model.update(_scenario_key.id, **{ 'overwritable': True })[1]
+        else:
+          scenario_model.update(_scenario_key.id, **{ 'overwritable': False })[1]
+  elif was_active and not is_active:
+    if _mode == intercept_mode.RECORD:
+      old_proxy_settings = __current_proxy_settings(context)
+      old_policy = __data_rule(old_proxy_settings).record_policy
+
+      if old_policy == record_policy.OVERWRITE:
+        scenario_key = data_rule.scenario_key
+        _scenario_key = __parse_scenario_key(scenario_key)
+
+        if _scenario_key:
+          scenario_model = ScenarioModel(new_settings)
+
+          # If policy is overwrite when recording, whenever intercept is disabled,
+          # set active scenario to not be overwritable
+          scenario_model.update(_scenario_key.id, **{ 'overwritable': False })[1]
+
 def handle_scenario_update(new_settings: Settings, context = None):
   new_scenario_key = __scenario_key(new_settings.proxy)
   if not new_scenario_key:
@@ -27,9 +65,10 @@ def handle_scenario_update(new_settings: Settings, context = None):
   old_proxy_settings = __current_proxy_settings(context)
   old_scenario_key = __scenario_key(old_proxy_settings)
 
-  data_rule = __data_rule(new_settings.proxy)
-  if data_rule.record_policy == record_policy.OVERWRITE:
-    if old_scenario_key != new_scenario_key:
+  if old_scenario_key != new_scenario_key:
+    data_rule = __data_rule(new_settings.proxy)
+
+    if data_rule.record_policy == record_policy.OVERWRITE:
       scenario_model = ScenarioModel(new_settings)
 
       _old_scenario_key = __parse_scenario_key(old_scenario_key)
@@ -75,13 +114,13 @@ def handle_policy_update(new_settings: Settings, context: Context = None):
     new_policy = data_rule.record_policy
     old_policy = old_data_rule.record_policy
 
-    if new_policy != old_policy and new_policy == record_policy.OVERWRITE:
+    if new_policy != old_policy and old_policy == record_policy.OVERWRITE:
       scenario_key = data_rule.scenario_key
       _scenario_key = __parse_scenario_key(scenario_key)
 
       if _scenario_key:
         scenario_model = ScenarioModel(new_settings)
-        scenario_model.update(_scenario_key.id, **{ 'overwritable': True })[1]
+        scenario_model.update(_scenario_key.id, **{ 'overwritable': False })[1]
 
 def __current_proxy_settings(context: Context = None):
   if context and context.get('current_proxy_settings'):

{stoobly-agent-0.29.1 → stoobly-agent-0.30.0}/stoobly_agent/app/cli/helpers/openapi_endpoint_adapter.py

@@ -1,13 +1,13 @@
 import copy
 from functools import reduce
 import itertools
-import pdb
 from pprint import pprint
 import re
 from typing import Dict, List, Union
 from urllib.parse import urlparse
 
 from openapi_core import Spec
+import yaml
 
 from stoobly_agent.lib.api.interfaces.endpoints import (
   Alias,
@@ -23,7 +23,14 @@ class OpenApiEndpointAdapter():
     return
 
   def adapt_from_file(self, file_path) -> List[EndpointShowResponse]:
-    spec = Spec.from_file_path(file_path)
+    spec = {}
+
+    with open(file_path, "r") as stream:
+      file_data: Dict = yaml.safe_load(stream)
+      if 'info' not in file_data:
+        self.__add_info(file_data)
+      spec = Spec.from_dict(file_data)
+
     return self.adapt(spec)
 
   def adapt(self, spec: Spec) -> List[EndpointShowResponse]:
@@ -33,7 +40,7 @@ class OpenApiEndpointAdapter():
     schemas = components.get("schemas", {})
     paths = spec.getkey('paths')
 
-    servers_spec = spec / "servers"
+    servers_spec = spec.get("servers")
     servers = self.__evaluate_servers(servers_spec)
 
     for _, server in enumerate(servers):
@@ -231,6 +238,15 @@ class OpenApiEndpointAdapter():
     
     return endpoints
 
+  def __add_info(self, file_data: Dict):
+    if 'info' in file_data:
+      return
+
+    file_data['info'] = {
+      'version': '0.0.1',
+      'title': ''
+    }
+
   def __get_most_recent_param(self, literal_params: dict):
     return list(literal_params)[-1] if literal_params else None
 

{stoobly-agent-0.29.1 → stoobly-agent-0.30.0}/stoobly_agent/app/cli/helpers/replay_facade.py

@@ -14,7 +14,7 @@ class ReplayCliOptions(TypedDict):
   assign: List[str]
   group_by: str
   host: str
-  lifecycle_hooks_script_path: str
+  lifecycle_hooks_path: str
   on_response: Callable
   project_key: str
   record: bool
@@ -54,7 +54,7 @@ class ReplayFacade():
       'before_replay': cli_options.get('before_replay'),
       'group_by': cli_options.get('group_by'),
       'host': cli_options.get('host'),
-      'lifecycle_hooks_script_path': cli_options.get('lifecycle_hooks_script_path'),
+      'lifecycle_hooks_path': cli_options.get('lifecycle_hooks_path'),
       'overwrite': cli_options.get('overwrite'),
       'request_origin': request_origin.CLI,
       'save': cli_options.get('save'),

{stoobly-agent-0.29.1 → stoobly-agent-0.30.0}/stoobly_agent/app/cli/intercept_cli.py

@@ -6,7 +6,7 @@ from stoobly_agent.app.settings import Settings
 from stoobly_agent.config.constants import mode, mock_policy, record_policy, replay_policy
 from stoobly_agent.lib.api.keys.project_key import ProjectKey
 
-from .helpers.handle_config_update_service import handle_policy_update
+from .helpers.handle_config_update_service import handle_intercept_active_update, handle_policy_update
 
 settings = Settings.instance()
 
@@ -43,6 +43,8 @@ def intercept(ctx):
 def enable(**kwargs):
     settings.proxy.intercept.active = True
 
+    handle_intercept_active_update(settings)
+
     settings.commit()
 
     print("Intercept enabled!")
@@ -55,6 +57,8 @@ def disable(**kwargs):
 
     settings.proxy.intercept.active = False
 
+    handle_intercept_active_update(settings)
+
     settings.commit()
 
     print("Intercept disabled!")
@@ -72,9 +76,14 @@ def configure(**kwargs):
         sys.exit(1)
 
     if kwargs['mode']:
-        settings.proxy.intercept.mode = kwargs['mode']
+        if settings.proxy.intercept.mode != kwargs['mode']:
+            if settings.proxy.intercept.active:
+                settings.proxy.intercept.active = False
+                handle_intercept_active_update(settings)
 
-        print(f"Updating intercept mode to {kwargs['mode']}")
+            settings.proxy.intercept.mode = kwargs['mode']
+
+            print(f"Updating intercept mode to {kwargs['mode']}")
 
     _mode = kwargs['mode'] or settings.proxy.intercept.mode 
 
@@ -98,9 +107,9 @@ def configure(**kwargs):
         elif active_mode == mode.TEST:
             data_rule.test_policy = kwargs['policy']
 
+        handle_policy_update(settings)
+
         print(f"Updating {_mode} policy to {kwargs['policy']}")
-    
-    handle_policy_update(settings)
 
     settings.commit()
 

{stoobly-agent-0.29.1 → stoobly-agent-0.30.0}/stoobly_agent/app/cli/request_cli.py

@@ -116,7 +116,7 @@ def list(**kwargs):
 @click.option('--format', type=click.Choice([BODY_FORMAT, JSON_FORMAT]), help='Format replay response.')
 @click.option('--host', help='Rewrite request host.')
 @ConditionalDecorator(lambda f: click.option('--group-by', help='Repeat for each alias name.')(f), is_remote)
-@click.option('--lifecycle-hooks-script-path', help='Path to lifecycle hooks script.')
+@click.option('--lifecycle-hooks-path', help='Path to lifecycle hooks script.')
 @click.option(
   '--log-level', default=logger.WARNING, type=click.Choice(log_levels), 
   help='''
@@ -215,7 +215,7 @@ if is_remote:
   @click.option('--format', type=click.Choice([BODY_FORMAT, JSON_FORMAT]), help='Format replay response.')
   @click.option('--group-by', help='Repeat for each alias name.')
   @click.option('--host', help='Rewrite request host.')
-  @click.option('--lifecycle-hooks-script-path', help='Path to lifecycle hooks script.')
+  @click.option('--lifecycle-hooks-path', help='Path to lifecycle hooks script.')
   @click.option(
       '--log-level', default=logger.WARNING, type=click.Choice(log_levels), 
       help='''

{stoobly-agent-0.29.1 → stoobly-agent-0.30.0}/stoobly_agent/app/cli/scenario_cli.py

@@ -65,7 +65,7 @@ def create(**kwargs):
 @click.option('--format', type=click.Choice([JSON_FORMAT]), help='Format replay response.')
 @ConditionalDecorator(lambda f: click.option('--group-by', help='Repeat for each alias name.')(f), is_remote)
 @click.option('--host', help='Rewrite request host.')
-@click.option('--lifecycle-hooks-script-path', help='Path to lifecycle hooks script.')
+@click.option('--lifecycle-hooks-path', help='Path to lifecycle hooks script.')
 @click.option(
     '--log-level', default=logger.WARNING, type=click.Choice([logger.DEBUG, logger.INFO, logger.WARNING, logger.ERROR]), 
     help='''
@@ -261,7 +261,7 @@ if is_remote:
     @click.option('--format', type=click.Choice([JSON_FORMAT]), help='Format replay response.')
     @click.option('--group-by', help='Repeat for each alias name.')
     @click.option('--host', help='Rewrite request host.')
-    @click.option('--lifecycle-hooks-script-path', help='Path to lifecycle hooks script.')
+    @click.option('--lifecycle-hooks-path', help='Path to lifecycle hooks script.')
     @click.option(
         '--log-level', default=logger.WARNING, type=click.Choice([logger.DEBUG, logger.INFO, logger.WARNING, logger.ERROR]), 
         help='''

{stoobly-agent-0.29.1 → stoobly-agent-0.30.0}/stoobly_agent/app/cli/types/endpoint.py

@@ -4,7 +4,7 @@ from stoobly_agent.app.models.types import OPENAPI_FORMAT
 
 class EndpointCreateCliOptions(TypedDict):
   format: Optional[Literal[f"{OPENAPI_FORMAT}"]]
-  lifecycle_hooks_script_path: Optional[str]
+  lifecycle_hooks_path: Optional[str]
   path: str
   project_key: Optional[str]
   scenario_key: Optional[str]

{stoobly-agent-0.29.1 → stoobly-agent-0.30.0}/stoobly_agent/app/models/adapters/raw_http_request_adapter.py

@@ -28,7 +28,7 @@ class RawHttpRequestAdapter():
 
         if decoded_key in self.headers:
           separator = HEADER_COOKIE_SEPARATOR if decoded_key.lower() == HEADER_COOKIE else HEADER_SEPARATOR
-          self.headers[decoded_key] = separator.join(self.headers[decoded_key], decoded_value)
+          self.headers[decoded_key] = separator.join([self.headers[decoded_key], decoded_value])
         else:
           self.headers[decoded_key] = decoded_value
 

{stoobly-agent-0.29.1 → stoobly-agent-0.30.0}/stoobly_agent/app/models/adapters/raw_http_response_adapter.py

@@ -35,7 +35,7 @@ class RawHttpResponseAdapter():
         decoded_key = decode(header_key)
         decoded_value = decode(header_value).strip()
         if decoded_key in self.headers:
-          self.headers[decoded_key] = HEADER_SEPARATOR.join(self.headers[decoded_key], decoded_value)
+          self.headers[decoded_key] = HEADER_SEPARATOR.join([self.headers[decoded_key], decoded_value])
         else:
           self.headers[decoded_key] = decoded_value
 

{stoobly-agent-0.29.1/stoobly_agent/app/proxy/record → stoobly-agent-0.30.0/stoobly_agent/app/proxy}/context.py

@@ -2,7 +2,8 @@ from mitmproxy.http import HTTPFlow as MitmproxyHTTPFlow
 
 from stoobly_agent.app.proxy.intercept_settings import InterceptSettings
 
-class RecordContext():
+class InterceptContext():
+
   def __init__(self, flow: MitmproxyHTTPFlow, intercept_settings: InterceptSettings):
     self.__flow = flow
     self.__intercept_settings = intercept_settings

{stoobly-agent-0.29.1 → stoobly-agent-0.30.0}/stoobly_agent/app/proxy/handle_mock_service.py

@@ -41,7 +41,7 @@ def handle_request_mock_generic(context: MockContext, **options: MockOptions):
         # Rewrite request with paramter rules for mock
         request: MitmproxyRequest = context.flow.request
         request_facade = MitmproxyRequestFacade(request)
-        request_facade.with_rewrite_rules(rewrite_rules).rewrite()
+        request_facade.with_parameter_rules(rewrite_rules).rewrite()
 
     __mock_hook(lifecycle_hooks.BEFORE_MOCK, context)
 

{stoobly-agent-0.29.1 → stoobly-agent-0.30.0}/stoobly_agent/app/proxy/handle_replay_service.py

@@ -28,10 +28,11 @@ def __replay_request(replay_context: ReplayContext):
     """
     intercept_settings: InterceptSettings = replay_context.intercept_settings
     rewrite_rules = intercept_settings.rewrite_rules
+
     if len(rewrite_rules) > 0:
         request: MitmproxyRequest = replay_context.flow.request
         request_facade = MitmproxyRequestFacade(request)
-        request_facade.with_rewrite_rules(rewrite_rules).rewrite()
+        request_facade.with_parameter_rules(rewrite_rules).with_url_rules(rewrite_rules).rewrite()
 
     __replay_hook(lifecycle_hooks.BEFORE_REPLAY, replay_context)
 

{stoobly-agent-0.29.1 → stoobly-agent-0.30.0}/stoobly_agent/app/proxy/intercept_handler.py

@@ -2,8 +2,9 @@ import os
 import pdb
 
 from mitmproxy.http import HTTPFlow as MitmproxyHTTPFlow
-from mitmproxy.http import Headers, Request as MitmproxyRequest, Response as MitmproxyResponse
+from mitmproxy.http import Headers, Request as MitmproxyRequest
 
+from stoobly_agent.app.proxy.context import InterceptContext
 from stoobly_agent.app.proxy.handle_mock_service import handle_request_mock
 from stoobly_agent.app.proxy.handle_replay_service import handle_request_replay, handle_response_replay
 from stoobly_agent.app.proxy.handle_record_service import handle_response_record
@@ -14,7 +15,7 @@ from stoobly_agent.app.proxy.replay.context import ReplayContext
 from stoobly_agent.app.proxy.record.context import RecordContext
 from stoobly_agent.app.proxy.utils.response_handler import bad_request
 from stoobly_agent.app.settings import Settings
-from stoobly_agent.config.constants import mode
+from stoobly_agent.config.constants import lifecycle_hooks, mode
 from stoobly_agent.lib.logger import Logger
 
 # Disable proxy settings in urllib
@@ -29,6 +30,8 @@ def request(flow: MitmproxyHTTPFlow):
 
     intercept_settings = InterceptSettings(Settings.instance(), request)
 
+    __intercept_hook(lifecycle_hooks.BEFORE_REQUEST, flow, intercept_settings)
+
     active_mode = intercept_settings.mode
     Logger.instance().debug(f"{LOG_ID}:ProxyMode: {active_mode}")
 
@@ -56,6 +59,8 @@ def response(flow: MitmproxyHTTPFlow):
     intercept_settings = InterceptSettings(Settings.instance(), request)
     intercept_settings.for_response()
 
+    __intercept_hook(lifecycle_hooks.BEFORE_RESPONSE, flow, intercept_settings)
+
     active_mode = intercept_settings.mode
 
     if active_mode == mode.RECORD:
@@ -90,4 +95,10 @@ def __patch_cookie(request: MitmproxyRequest):
         
 def __combine_header(headers: Headers, header_name: str, delimitter: str):
     values = headers.get_all(header_name)
-    headers[header_name] = delimitter.join(values)
+    headers[header_name] = delimitter.join(values)
+
+def __intercept_hook(hook: str, flow: MitmproxyHTTPFlow, intercept_settings: InterceptSettings):
+    lifecycle_hooks_module = intercept_settings.lifecycle_hooks
+
+    if hook in lifecycle_hooks_module:
+        lifecycle_hooks_module[hook](InterceptContext(flow, intercept_settings))

{stoobly-agent-0.29.1 → stoobly-agent-0.30.0}/stoobly_agent/app/proxy/intercept_settings.py

@@ -10,8 +10,8 @@ from stoobly_agent.app.settings.constants import firewall_action, intercept_mode
 from stoobly_agent.app.settings.rewrite_rule import RewriteRule
 from stoobly_agent.app.settings.firewall_rule import FirewallRule
 from stoobly_agent.app.settings import Settings
-from stoobly_agent.app.settings.types import IgnoreRule, MatchRule, RedactRule
-from stoobly_agent.config.constants import custom_headers, mode, request_origin, test_filter
+from stoobly_agent.app.settings.types import IgnoreRule, MatchRule
+from stoobly_agent.config.constants import custom_headers, env_vars, mode, request_origin, test_filter
 from stoobly_agent.lib.api.keys.project_key import InvalidProjectKey, ProjectKey
 from stoobly_agent.lib.logger import Logger
 
@@ -62,9 +62,12 @@ class InterceptSettings:
     return custom_headers.PROXY_MODE in self.__headers
 
   @property
-  def lifecycle_hooks_script_path(self):
-    if self.__headers and custom_headers.LIFECYCLE_HOOKS_SCRIPT_PATH in self.__headers:
-      return self.__headers[custom_headers.LIFECYCLE_HOOKS_SCRIPT_PATH]
+  def lifecycle_hooks_path(self):
+    if self.__headers and custom_headers.LIFECYCLE_HOOKS_PATH in self.__headers:
+      return self.__headers[custom_headers.LIFECYCLE_HOOKS_PATH]
+
+    if os.environ.get(env_vars.AGENT_LIFECYCLE_HOOKS_PATH):
+      return os.environ[env_vars.AGENT_LIFECYCLE_HOOKS_PATH] 
 
   @property
   def lifecycle_hooks(self):
@@ -151,11 +154,11 @@ class InterceptSettings:
     return self.__select_rewrite_rules(_mode)
 
   @property
-  def record_rewrite_rules(self) -> List[RedactRule]:
+  def record_rewrite_rules(self) -> List[RewriteRule]:
     return self.__select_rewrite_rules(mode.RECORD)
 
   @property
-  def mock_rewrite_rules(self) -> List[RedactRule]:
+  def mock_rewrite_rules(self) -> List[RewriteRule]:
     return self.__select_rewrite_rules(mode.MOCK)
 
   @property
@@ -207,21 +210,18 @@ class InterceptSettings:
 
     # Filter only parameters matching active intercept mode
     for rewrite_rule in self.__rewrite_rules:
-      parameter_rules = self.__select_parameter_rules(rewrite_rule, mode)
+      # If url rule applies, then update .url_rules with url_rule
+      url_rules = self.__select_url_rules(rewrite_rule)
 
-      # If no parameters rules were found, then this filter rule is not applied
-      if len(parameter_rules) == 0:
-        continue
-
-      # Build a new RewriteRule object contain only parameter rules matching intercept mode
-      rewrite_rule = RewriteRule({
-        'methods': rewrite_rule.methods,
-        'pattern': rewrite_rule.pattern,
-        'parameters_rules': [], # Has to be dict form, manually set it
-      })
-      rewrite_rule.parameter_rules = parameter_rules
+      # If parameters rules apply, then update .parameter_rules with applicable parameter_rules
+      parameter_rules = self.__select_parameter_rules(rewrite_rule, mode)
 
-      rules.append(rewrite_rule)
+      if len(url_rules) > 0 or len(parameter_rules) > 0:
+        # Build a new RewriteRule object contain only parameter rules matching intercept mode
+        rewrite_rule = RewriteRule(rewrite_rule.to_dict())
+        rewrite_rule.url_rules = url_rules
+        rewrite_rule.parameter_rules = parameter_rules
+        rules.append(rewrite_rule)
 
     return rules
 
@@ -232,8 +232,15 @@ class InterceptSettings:
       rewrite_rule.parameter_rules or []
     ))
 
+  def __select_url_rules(self, rewrite_rule: RewriteRule, mode = None):
+    mode = mode or self.mode
+    return list(filter(
+      lambda url: mode in url.modes,
+      rewrite_rule.url_rules or []
+    ))
+
   def __initialize_lifecycle_hooks(self):
-    script_path = self.lifecycle_hooks_script_path
+    script_path = self.lifecycle_hooks_path
 
     if not script_path:
         return