stoobly-agent 0.29.1__py3-none-any.whl → 0.30.0__py3-none-any.whl

stoobly_agent/app/proxy/mitmproxy/request_facade.py

@@ -8,7 +8,7 @@ from typing import Callable, List
 from urllib.parse import urlparse
 
 from stoobly_agent.app.settings.constants import request_component
-from stoobly_agent.app.settings.rewrite_rule import RewriteRule, ParameterRule
+from stoobly_agent.app.settings.rewrite_rule import ParameterRule, RewriteRule, UrlRule 
 from stoobly_agent.config.constants import custom_headers
 from stoobly_agent.lib.logger import Logger, bcolors
 from stoobly_agent.lib.utils import jmespath
@@ -29,8 +29,8 @@ class MitmproxyRequestFacade(Request):
         self.request = request
         self.uri = urlparse(self.request.url)
 
-        self.__redact_rules: List[ParameterRule] = []
-        self.__rewrite_rules: List[ParameterRule] = []
+        self.__url_rules: List[UrlRule] = []
+        self.__parameter_rules: List[ParameterRule] = []
 
         self.__body = MitmproxyRequestBodyFacade(request)
 
@@ -99,41 +99,42 @@ class MitmproxyRequestFacade(Request):
         return self.request.port
 
     @property
-    def redact_rules(self) -> List[ParameterRule]:
-        return self.__redact_rules
+    def url_rules(self) -> List[ParameterRule]:
+        return self.__url_rules
     
     @property
-    def rewrite_rules(self) -> List[ParameterRule]:
-        return self.__rewrite_rules
+    def parameter_rules(self) -> List[ParameterRule]:
+        return self.__parameter_rules
 
-    def with_rewrite_rules(self, rules: List[RewriteRule]):
+    def with_parameter_rules(self, rules: List[RewriteRule]):
         if type(rules) == list:
-            self.__rewrite_rules = self.select_parameter_rules(rules)
+            self.__parameter_rules = self.select_parameter_rules(rules)
         return self 
 
-    def with_redact_rules(self, rules: List[RewriteRule]):
+    def with_url_rules(self, rules: List[RewriteRule]):
         if type(rules) == list:
-            self.__redact_rules = self.select_parameter_rules(rules)
+            self.__url_rules = self.select_url_rules(rules)
         return self
 
-    def redact(self):
-        redacts = self.__redact_rules
-        if len(redacts) != 0:
-            self.__redact_headers(redacts)
-            self.__redact_query(redacts)
-            self.__redact_content(redacts)
-
     def rewrite(self):
-        rewrites = self.__rewrite_rules
+        rewrites = self.__parameter_rules
 
         if len(rewrites) != 0:
             self.__rewrite_headers(rewrites)
             self.__rewrite_query(rewrites)
             self.__rewrite_content(rewrites)
 
+        rewrites = self.__url_rules
+
+        if len(rewrites):
+            self.__rewrite_url(rewrites)
+
+    # Find all the rules that match request url and method
+    def select_rewrite_rules(self, rules: List[RewriteRule]) -> List[RewriteRule]:
+        return list(filter(self.__is_rewrite_rule_selected, rules or []))
+
     def select_parameter_rules(self, rules: List[RewriteRule]) -> List[ParameterRule]:
-        # Find all the rules that match request url and method
-        _rules = list(filter(self.__is_parameter_rule_selected, rules or []))
+        _rules = self.select_rewrite_rules(rules)
         
         if len(_rules) == 0:
             return []
@@ -142,7 +143,17 @@ class MitmproxyRequestFacade(Request):
 
         return [item for sublist in parameter_rules for item in sublist] # flatten list
 
-    def __is_parameter_rule_selected(self, rewrite_rule: RewriteRule):
+    def select_url_rules(self, rules: List[RewriteRule]) -> List[UrlRule]:
+        _rules = self.select_rewrite_rules(rules)
+
+        if len(_rules) == 0:
+            return []
+
+        url_rules = list(map(lambda rule: rule.url_rules, _rules))
+
+        return [item for sublist in url_rules for item in sublist] # flatten list
+
+    def __is_rewrite_rule_selected(self, rewrite_rule: RewriteRule):
         pattern = rewrite_rule.pattern
 
         try:
@@ -168,6 +179,14 @@ class MitmproxyRequestFacade(Request):
         Logger.instance().info(f"{bcolors.OKCYAN}Rewriting {rewrite.type.lower()}{bcolors.ENDC} {rewrite.name} => {rewrite.value}")
         return rewrite.value
 
+    def __rewrite_url(self, rewrites: List[UrlRule]):
+        for rewrite in rewrites:
+            if rewrite.host:
+                self.request.host = rewrite.host
+            
+            if rewrite.port:
+                self.request.port = rewrite.port
+
     def __rewrite_headers(self, rewrites: List[ParameterRule]):
         self.__apply_headers(rewrites, self.__rewrite_handler)
 
@@ -177,19 +196,6 @@ class MitmproxyRequestFacade(Request):
     def __rewrite_content(self, rewrites: List[ParameterRule]):
         self.__apply_content(rewrites, self.__rewrite_handler)
 
-    def __redact_handler(self, rewrite: ParameterRule) -> str:
-        Logger.instance().debug(f"{bcolors.OKCYAN}Redacting{bcolors.ENDC} {rewrite.name}")
-        return '[REDACTED]'
-
-    def __redact_headers(self, redacts: List[ParameterRule]):
-        self.__apply_headers(redacts, self.__redact_handler)
-
-    def __redact_query(self, redacts: List[ParameterRule]):
-        self.__apply_queries(redacts, self.__redact_handler)
-
-    def __redact_content(self, redacts: List[ParameterRule]):
-        self.__apply_content(redacts, self.__redact_handler)
-
     def __apply_headers(self, rewrites: List[ParameterRule], handler: Callable):
         rewrites = list(filter(lambda rewrite: rewrite.type == request_component.HEADER, rewrites))
         self.__apply_rewrites(self.request.headers, rewrites, handler)

stoobly_agent/app/proxy/mitmproxy/response_facade.py

@@ -11,7 +11,6 @@ class MitmproxyResponseFacade(Response):
         self.content = response.raw_content
 
         self.rewrite_rules = []
-        self.rewrite_rules = []
 
     @property
     def code(self):
@@ -42,22 +41,12 @@ class MitmproxyResponseFacade(Response):
         self.response.headers['content-length'] = str(len(self.content))
 
     
-    def with_redact_rules(self, rules: RewriteRule):
-        if type(rules) == list:
-            self.rewrite_rules = rules
-
-        return self
-
     def with_rewrite_rules(self, rules: RewriteRule):
         if type(rules) == list:
             self.rewrite_rules = rules
 
         return self
 
-    # TODO
-    def redact(self):
-        pass
-    
     # TODO
     def rewrite(self):
         pass

stoobly_agent/app/proxy/mock/context.py

@@ -5,10 +5,11 @@ from requests import Response
 
 from stoobly_agent.app.proxy.intercept_settings import InterceptSettings
 
-class MockContext():
+from ..context import InterceptContext
+
+class MockContext(InterceptContext):
   def __init__(self, flow: MitmproxyHTTPFlow, intercept_settings: InterceptSettings):
-    self.__flow = flow
-    self.__intercept_settings = intercept_settings
+    super().__init__(flow, intercept_settings)
 
     self.__start_time = time.time()
     self.__end_time = None
@@ -18,14 +19,6 @@ class MockContext():
   def end_time(self):
     return self.__end_time
 
-  @property
-  def flow(self):
-    return self.__flow
-
-  @property
-  def intercept_settings(self):
-    return self.__intercept_settings
-
   @property
   def response(self):
     return self.__response

stoobly_agent/app/proxy/record/context.py

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

stoobly_agent/app/proxy/record/join_request_service.py

@@ -26,7 +26,7 @@ def join_rewritten_request(flow: MitmproxyHTTPFlow, intercept_settings: Intercep
     response = MitmproxyResponseFacade(flow.response)
     rewrite_rules = intercept_settings.record_rewrite_rules
 
-    request.with_rewrite_rules(rewrite_rules).rewrite()
+    request.with_parameter_rules(rewrite_rules).rewrite()
     response.with_rewrite_rules(rewrite_rules).rewrite()
 
     return join_request(request, response, intercept_settings)

stoobly_agent/app/proxy/replay/context.py

@@ -5,10 +5,11 @@ from requests import Response
 
 from stoobly_agent.app.proxy.intercept_settings import InterceptSettings
 
-class ReplayContext():
+from ..context import InterceptContext
+
+class ReplayContext(InterceptContext):
   def __init__(self, flow: MitmproxyHTTPFlow, intercept_settings: InterceptSettings):
-    self.__flow = flow
-    self.__intercept_settings = intercept_settings
+    super().__init__(flow, intercept_settings)
 
     self.__start_time = time.time()
     self.__end_time = None
@@ -18,14 +19,6 @@ class ReplayContext():
   def end_time(self):
     return self.__end_time
 
-  @property
-  def flow(self):
-    return self.__flow
-
-  @property
-  def intercept_settings(self):
-    return self.__intercept_settings
-
   @property
   def response(self):
     return self.__response

stoobly_agent/app/proxy/replay/replay_request_service.py

@@ -21,7 +21,7 @@ class ReplayRequestOptions(TypedDict):
   alias_resolve_strategy: alias_resolve_strategy.AliasResolveStrategy
   group_by: str
   host: str
-  lifecycle_hooks_script_path: str
+  lifecycle_hooks_path: str
   mode: Union[mode.MOCK, mode.RECORD, mode.TEST, None]
   before_replay: Union[Callable[[ReplayContext], None], None]
   after_replay: Union[Callable[[ReplayContext], Union[requests.Response, None]], None]
@@ -54,8 +54,8 @@ def replay(context: ReplayContext, options: ReplayRequestOptions) -> requests.Re
   if options.get('alias_resolve_strategy'):
     headers[custom_headers.ALIAS_RESOLVE_STRATEGY] = options['alias_resolve_strategy']
 
-  if options.get('lifecycle_hooks_script_path'):
-    __handle_lifecycle_hooks_script_path(options['lifecycle_hooks_script_path'], headers) 
+  if options.get('lifecycle_hooks_path'):
+    __handle_lifecycle_hooks_path(options['lifecycle_hooks_path'], headers) 
 
   if options.get('mode'):
     __handle_mode_option(options['mode'], request, headers)
@@ -137,14 +137,14 @@ def replay(context: ReplayContext, options: ReplayRequestOptions) -> requests.Re
 
   return res
 
-def __handle_lifecycle_hooks_script_path(script_path, headers):
+def __handle_lifecycle_hooks_path(script_path, headers):
   if not script_path:
     return
 
   if not os.path.isabs(script_path):
     script_path = os.path.join(os.path.abspath('.'), script_path)
 
-  headers[custom_headers.LIFECYCLE_HOOKS_SCRIPT_PATH] = script_path
+  headers[custom_headers.LIFECYCLE_HOOKS_PATH] = script_path
 
 def __handle_mode_option(_mode, request: Request, headers):
   headers[custom_headers.PROXY_MODE] = _mode

stoobly_agent/app/proxy/replay/rewrite_params_service.py

@@ -1,6 +1,6 @@
 import pdb
 
-from typing import Callable, Dict, List, TypedDict, Union
+from typing import Dict, List, Union
 
 from stoobly_agent.app.proxy.replay.alias_context import AliasContext
 from stoobly_agent.app.proxy.replay.alias_resolver import AliasResolver

stoobly_agent/app/proxy/run.py

@@ -113,6 +113,9 @@ def __filter_options(options):
     if 'headless' in options:
         del options['headless']
 
+    if 'lifecycle_hooks_path' in options:
+        del options['lifecycle_hooks_path']
+
     if 'ui_host' in options:
         del options['ui_host']
 

stoobly_agent/app/proxy/test/context.py

@@ -74,8 +74,8 @@ class TestContext(TestContextABC):
     return self.__intercept_settings.lifecycle_hooks
 
   @property
-  def lifecycle_hooks_script_path(self):
-    return self.__intercept_settings.lifecycle_hooks_script_path
+  def lifecycle_hooks_path(self):
+    return self.__intercept_settings.lifecycle_hooks_path
 
   @property
   def end_time(self):

stoobly_agent/app/proxy/test/context_abc.py

@@ -35,7 +35,7 @@ class TestContextABC(abc.ABC):
 
   @property
   @abc.abstractmethod
-  def lifecycle_hooks_script_path(self):
+  def lifecycle_hooks_path(self):
     pass
 
   @property

stoobly_agent/app/proxy/test/matchers/custom.py

@@ -4,7 +4,7 @@ def matches(context: TestContext):
     lifecycle_hooks = context.lifecycle_hooks 
 
     if not 'handle_test' in lifecycle_hooks:
-        return False, f"Expected function 'handle_test' to be defined in {context.lifecycle_hooks_script_path}"
+        return False, f"Expected function 'handle_test' to be defined in {context.lifecycle_hooks_path}"
 
     try:
         status, log = lifecycle_hooks['handle_test'](context)

stoobly_agent/app/settings/rewrite_rule.py

@@ -2,6 +2,7 @@ import pdb
 
 from .parameter_rule import ParameterRule
 from .types.proxy_settings import RewriteRule as IRewriteRule
+from .url_rule import UrlRule
 
 class RewriteRule:
 
@@ -14,6 +15,9 @@ class RewriteRule:
     self.__raw_parameter_rules = self.__rewrite_rule.get('parameter_rules') or []
     self.__parameter_rules = list(map(lambda rule: ParameterRule(rule), self.__raw_parameter_rules))
 
+    self.__raw_url_rules = self.__rewrite_rule.get('url_rules') or []
+    self.__url_rules = list(map(lambda rule: UrlRule(rule), self.__raw_url_rules))
+
   @property 
   def methods(self):
     return self.__methods
@@ -30,9 +34,18 @@ class RewriteRule:
   def parameter_rules(self, v):
     self.__parameter_rules = v
 
+  @property
+  def url_rules(self):
+    return self.__url_rules
+
+  @url_rules.setter
+  def url_rules(self, v):
+    self.__url_rules = v
+
   def to_dict(self):
     return {
       'methods': self.__methods,
       'pattern': self.__pattern,
       'parameter_rules': list(map(lambda parameter: parameter.to_dict(), self.__parameter_rules)),
+      'url_rules': list(map(lambda url: url.to_dict(), self.__url_rules)),
     }

stoobly_agent/app/settings/types/__init__.py

@@ -23,11 +23,6 @@ class IgnoreRule(TypedDict):
     method: str
     pattern: str
 
-class RedactRule(TypedDict):
-    redacts: List[Redact]
-    method: str
-    pattern: str
-
 class RewriteRule(TypedDict):
     rewrites: List[Rewrite]
     method: str
@@ -97,4 +92,4 @@ Component = {
     'QueryParam': 'Query Param',
 }
 IProjectModeSettings = Union[IProjectMockSettings, IProjectRecordSettings, IProjectTestSettings]
-Rule = Union[IgnoreRule, RedactRule, RewriteRule]
+Rule = Union[IgnoreRule, RewriteRule]

stoobly_agent/app/settings/types/proxy_settings.py

@@ -15,6 +15,11 @@ class ParameterRule(TypedDict):
   type: str
   value: str
 
+class UrlRule(TypedDict):
+  host: str
+  modes: List[Mode]
+  port: str
+
 class DataRules(TypedDict):
   mock_policy: MockPolicy
   record_policy: RecordPolicy
@@ -30,6 +35,7 @@ class RewriteRule(TypedDict):
   methods: List[Method]
   pattern: str
   parameter_rules: List[ParameterRule]
+  url_rule: List[UrlRule]
 
 class FirewallRule(TypedDict):
   action: List[FirewallAction]

stoobly_agent/app/settings/url_rule.py

@@ -0,0 +1,33 @@
+from typing import List
+
+from .types.proxy_settings import UrlRule as IUrlRule
+
+class UrlRule:
+
+  def __init__(self, url_rule: IUrlRule):
+    self.update(url_rule)
+
+  @property
+  def host(self):
+    return self.__host
+
+  @property 
+  def modes(self):
+    return self.__modes
+
+  @property
+  def port(self):
+    return self.__port
+
+  def update(self, url_rule: IUrlRule):
+    self.__url_rule = url_rule
+    self.__host = self.__url_rule.get('host')
+    self.__modes = self.__url_rule.get('modes')
+    self.__port = self.__url_rule.get('port')
+
+  def to_dict(self):
+    return {
+      'host': self.__host,
+      'modes': self.__modes,
+      'port': self.__port,
+    }

stoobly_agent/cli.py

@@ -87,6 +87,7 @@ def init(**kwargs):
 @click.option('--log-level', default=logger.INFO, type=click.Choice([logger.DEBUG, logger.INFO, logger.WARNING, logger.ERROR]), help='''
     Log levels can be "debug", "info", "warning", or "error"
 ''')
+@click.option('--lifecycle-hooks-path', help='Path to lifecycle hooks script.')
 @click.option('--modify-headers', multiple=True, help='''
   Header modify pattern of the form "[/flow-filter]/header-name/[@]header-value", where the separator can be any character. The @ allows to provide a file path that is used to read the header value string.
   An empty header-value removes existing header-name headers. May be passed multiple times.
@@ -107,6 +108,9 @@ def run(**kwargs):
 
     os.environ[env_vars.AGENT_PROXY_URL] = f"http://{kwargs['proxy_host']}:{kwargs['proxy_port']}"
 
+    if kwargs.get('lifecycle_hooks_path'):
+      os.environ[env_vars.AGENT_LIFECYCLE_HOOKS_PATH] = kwargs['lifecycle_hooks_path']
+
     # Observe config for changes
     Settings.instance().watch()
 

stoobly_agent/config/constants/custom_headers.py

@@ -3,7 +3,7 @@ MOCK_POLICY = 'X-Mock-Policy'
 MOCK_REQUEST_ID = 'X-Stoobly-Request-Id'
 MOCK_REQUEST_ENDPOINT_ID = 'X-Stoobly-Request-Endpoint-Id'
 DO_PROXY = 'X-Do-Proxy'
-LIFECYCLE_HOOKS_SCRIPT_PATH = 'X-Stoobly-Lifecycle-Hooks-Script-Path'
+LIFECYCLE_HOOKS_PATH = 'X-Stoobly-Lifecycle-Hooks-Path'
 PROJECT_KEY = 'X-Project-Key'
 PROXY_MODE = 'X-Proxy-Mode'
 RECORD_POLICY = 'X-Record-Policy'

stoobly_agent/config/constants/env_vars.py

@@ -5,6 +5,7 @@ AGENT_SELF_INTERCEPT_ENABLED = 'STOOBLY_AGENT_SELF_INTERCEPT_ENABLED'
 AGENT_INCLUDE_PATTERNS = 'STOOBLY_AGENT_INCLUDE_PATTERNS'
 AGENT_IS_HEADLESS = 'STOOBLY_AGENT_IS_HEADLESS'
 AGENT_EXCLUDE_PATTERNS = 'STOOBLY_AGENT_EXCLUDE_PATTERNS'
+AGENT_LIFECYCLE_HOOKS_PATH = 'STOOBLY_AGENT_LIFECYCLE_HOOKS_PATH'
 AGENT_POLICY = 'STOOBLY_AGENT_POLICY'
 AGENT_PROJECT_KEY = 'STOOBLY_AGENT_PROJECT_KEY'
 AGENT_PROXY_URL = 'STOOBLY_AGENT_PROXY_URL'

stoobly_agent/config/constants/lifecycle_hooks.py

@@ -5,6 +5,8 @@ AFTER_TEST = 'handle_after_test'
 BEFORE_MOCK = 'handle_before_mock'
 BEFORE_RECORD = 'handle_before_record'
 BEFORE_REPLAY = 'handle_before_replay'
+BEFORE_REQUEST = 'handle_before_request'
+BEFORE_RESPONSE = 'handle_before_response'
 BEFORE_TEST = 'handle_before_test'
 ON_LENGTH_MATCH_ERROR = 'handle_length_match_error'
 ON_PARAM_NAME_EXISTS_ERROR = 'handle_param_name_exists_error'

stoobly_agent/config/schema.yml

@@ -48,6 +48,12 @@ rewrite_rule:
   methods: list(include('method'))
   pattern: str()
   parameter_rules: list(include('parameter_rule'))
+  url_rules: list(include('url_rule'))
+  port: str()
+url_rule:
+  host: str()
+  modes: list(include('mode'), min=1, max=4)
+  port: str()
 firewall_rule:
   action: include('firewall_action')
   methods: list(include('method'))
@@ -57,4 +63,4 @@ match_rule:
   components: include('component')
   methods: list(include('method'))
   modes: list(include('mode'), min=1, max=4)
-  pattern: str()
+  pattern: str()

stoobly_agent/lib/orm/request.py

@@ -155,11 +155,15 @@ def handle_deleting(request):
     replayed_response.delete()
 
 def handle_deleted(request):
-  scenario = request.scenario
+  request_before = request.get_original()
+
+  if request_before.get('scenario_id'):
+    # request.scenario returns a cached version of the scenario
+    scenario = Scenario.find(request_before['scenario_id'])
 
-  if scenario:
-    scenario.requests_count -= 1
-    scenario.save()
+    if scenario:
+      scenario.requests_count -= 1
+      scenario.save()
 
 Request.saved(handle_saved)
 Request.saving(handle_saving)