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

stoobly_agent/__init__.py

@@ -1,2 +1,2 @@
 COMMAND = 'stoobly-agent'
-VERSION = '1.4.2'
+VERSION = '1.5.1'

stoobly_agent/app/cli/helpers/handle_config_update_service.py

@@ -58,9 +58,9 @@ def handle_intercept_active_update(new_settings: Settings, context: Context = No
           scenario_model.update(_scenario_key.id, **{ 'overwritable': False })[1]
     elif _mode == intercept_mode.MOCK:
       # When mock is stopped, clear request access counts
-      from stoobly_agent.app.models.factories.resource.local_db.helpers.tiebreak_scenario_request import reset
+      from stoobly_agent.app.models.factories.resource.local_db.helpers.tiebreak_scenario_request import reset_sessions
 
-      reset()
+      reset_sessions()
 
 def handle_scenario_update(new_settings: Settings, context = None):
   new_scenario_key = __scenario_key(new_settings.proxy)

stoobly_agent/app/cli/helpers/handle_mock_service.py

@@ -6,9 +6,13 @@ from stoobly_agent.app.proxy.record.response_string import ResponseString
 
 RAW_FORMAT = 'raw'
 
-def print_raw_response(response: requests.Response):
+def print_raw_response(response: requests.Response, file_path = None):
   mitmproxy_response = PythonResponseAdapterFactory(response).mitmproxy_response()
   facade = MitmproxyResponseFacade(mitmproxy_response)
   response_string = ResponseString(facade, None)
 
-  print(response_string.get().decode(), end="")
+  if not file_path:
+    print(response_string.get().decode(), end="")
+  else:
+    with open(file_path, 'w') as fp:
+      fp.write(response_string.get().decode()) 

stoobly_agent/app/cli/helpers/request_facade.py

@@ -53,10 +53,14 @@ class RequestFacade(ReplayFacade):
   def replay(self, request_key: str, cli_options: ReplayCliOptions):
     replay_context = self.__build_replay_context(request_key)
     replay_options = {
-      'mode': mode.RECORD if cli_options.get('record') else mode.REPLAY,
+      'mode': mode.REPLAY,
       **self.__common_replay_options(request_key),
       **self.common_replay_cli_options(cli_options)
     }
+
+    if cli_options.get('record'):
+      replay_options['response_mode'] = mode.RECORD
+
     trace_context = replay_options.get('trace_context')
 
     return self.__replay(replay_context, trace_context, replay_options)

stoobly_agent/app/cli/scaffold/constants.py

@@ -14,8 +14,8 @@ CONFIG_FILE = '.config.yml'
 CONTEXT_DIR_ENV = 'CONTEXT_DIR'
 DOCKER_NAMESPACE = 'docker'
 ENV_FILE = '.env'
-FIXTURES_FOLDER_NAME = 'fixtures'
 NAMESERVERS_FILE = '.nameservers'
+PUBLIC_FOLDER_NAME = 'public'
 SERVICE_DETACHED = '${SERVICE_DETACHED}'
 SERVICE_DETACHED_ENV = 'SERVICE_DETACHED'
 SERVICE_HOSTNAME = '${SERVICE_HOSTNAME}'

stoobly_agent/app/cli/scaffold/docker/workflow/mock_decorator.py

@@ -26,6 +26,7 @@ class MockDecorator():
       '--lifecycle-hooks-path', 'lifecycle_hooks.py',
       '--proxy-mode', SERVICE_PROXY_MODE,
       '--proxy-port', f"{SERVICE_PORT}",
+      '--public-directory-path', 'public',
       '--response-fixtures-path', 'fixtures.yml',
       '--ssl-insecure'
     ]

stoobly_agent/app/cli/scaffold/service_workflow_validate_command.py

@@ -10,7 +10,7 @@ import yaml
 from docker.models.containers import Container
 
 from stoobly_agent.app.cli.scaffold.constants import (
-  FIXTURES_FOLDER_NAME,
+  PUBLIC_FOLDER_NAME,
   STOOBLY_DATA_DIR,
   VIRTUAL_HOST_ENV,
   VIRTUAL_PORT_ENV,
@@ -40,8 +40,8 @@ class ServiceWorkflowValidateCommand(ServiceCommand, ValidateCommand):
     )
 
   @property
-  def fixtures_dir_path(self):
-    return os.path.join(self.workflow_path, FIXTURES_FOLDER_NAME)
+  def public_dir_path(self):
+    return os.path.join(self.workflow_path, PUBLIC_FOLDER_NAME)
   
   @property
   def workflow_path(self):
@@ -124,14 +124,14 @@ class ServiceWorkflowValidateCommand(ServiceCommand, ValidateCommand):
     if ('200 OK' not in logs) and ('499' not in logs):
       raise ScaffoldValidateException(f"Error reaching {url} from inside Docker network")
 
-  # Check fixtures folder mounted into container
-  def validate_fixtures_folder(self, container: Container):
+  # Check public folder exists in container
+  def validate_public_folder(self, container: Container):
     
     if self.workflow_name == WORKFLOW_RECORD_TYPE:
-      print(f"Skipping validating fixtures folder in workflow: {self.workflow_name}, container: {container.name}")
+      print(f"Skipping validating public folder in workflow: {self.workflow_name}, container: {container.name}")
       return
 
-    print(f"Validating fixtures folder in container: {container.name}")
+    print(f"Validating public folder in container: {container.name}")
 
     data_dir_mounted = False
     volume_mounts = container.attrs['Mounts']
@@ -145,21 +145,21 @@ class ServiceWorkflowValidateCommand(ServiceCommand, ValidateCommand):
 
     # Only the running proxy containers will be checkable
     if container.status == 'exited':
-      print(f"Skipping validating fixtures folder contents because container is exited: {container.name}")
+      print(f"Skipping validating public folder contents because container is exited: {container.name}")
       return
 
-    # Check contents of fixtures folder to confirm it's shared
-    fixtures_folder_path = f"{FIXTURES_FOLDER_NAME}"
-    exec_result = container.exec_run(f"ls -A {fixtures_folder_path}")
+    # Check contents of public folder to confirm it's shared
+    public_folder_path = f"{PUBLIC_FOLDER_NAME}"
+    exec_result = container.exec_run(f"ls -A {public_folder_path}")
     output = exec_result.output
 
-    fixtures_folder_contents_container = output.decode('ascii').split('\n')
-    if fixtures_folder_contents_container[-1] == '':
-      fixtures_folder_contents_container.pop()
-    fixtures_folder_contents_scaffold = os.listdir(self.fixtures_dir_path)
+    public_folder_contents_container = output.decode('ascii').split('\n')
+    if public_folder_contents_container[-1] == '':
+      public_folder_contents_container.pop()
+    public_folder_contents_scaffold = os.listdir(self.public_dir_path)
 
-    if Counter(fixtures_folder_contents_container) != Counter(fixtures_folder_contents_scaffold):
-      raise ScaffoldValidateException(f"Fixtures was not mounted properly, expected {self.fixtures_dir_path} to exist in container path {fixtures_folder_path}")
+    if Counter(public_folder_contents_container) != Counter(public_folder_contents_scaffold):
+      raise ScaffoldValidateException(f"public folder was not mounted properly, expected {self.public_dir_path} to exist in container path {public_folder_path}")
 
   # Note: might not need this if the hostname is reachable and working
   def proxy_environment_variables_exist(self, container: Container) -> None:
@@ -191,7 +191,7 @@ class ServiceWorkflowValidateCommand(ServiceCommand, ValidateCommand):
       raise ScaffoldValidateException(f"Container attributes are missing for: {container.name}")
 
     if not self.service_config.detached:
-      self.validate_fixtures_folder(service_proxy_container)
+      self.validate_public_folder(service_proxy_container)
 
     self.proxy_environment_variables_exist(service_proxy_container)
 
@@ -215,7 +215,7 @@ class ServiceWorkflowValidateCommand(ServiceCommand, ValidateCommand):
 
     # Service init containers have a mounted dist folder unlike the core init container
     init_container = self.docker_client.containers.get(self.service_docker_compose.init_container_name)
-    self.validate_fixtures_folder(init_container)
+    self.validate_public_folder(init_container)
 
     if self.service_config.hostname:
       service_proxy_container = self.docker_client.containers.get(self.service_docker_compose.proxy_container_name)

stoobly_agent/app/cli/scaffold/templates/app/.Dockerfile.context

@@ -1,4 +1,4 @@
-FROM stoobly/agent:1.4
+FROM stoobly/agent:1.5
 
 ARG USER_ID
 

stoobly_agent/app/cli/scaffold/templates/constants.py

@@ -18,12 +18,12 @@ CUSTOM_FIXTURES = 'fixtures.yml'
 CUSTOM_LIFECYCLE_HOOKS = os.path.join('lifecycle_hooks.py')
 MAINTAINED_CONFIGURE = os.path.join('bin', '.configure')
 MAINTAINED_INIT = os.path.join('bin', '.init')
-MAINTAINED_FIXTURES = os.path.join('fixtures', '.keep')
+MAINTAINED_PUBLIC = os.path.join('public', '.gitignore')
 
 MOCK_WORKFLOW_MAINTAINED_FILES = [
   MAINTAINED_CONFIGURE,
   MAINTAINED_INIT,
-  MAINTAINED_FIXTURES
+  MAINTAINED_PUBLIC
 ]
 
 MOCK_WORKFLOW_CUSTOM_FILES = [
@@ -49,7 +49,7 @@ RECORD_WORKFLOW_CUSTOM_FILES = [
 TEST_WORKFLOW_MAINTAINED_FILES = [
   MAINTAINED_CONFIGURE,
   MAINTAINED_INIT,
-  MAINTAINED_FIXTURES
+  MAINTAINED_PUBLIC
 ]
 
 TEST_WORKFLOW_CUSTOM_FILES = [

stoobly_agent/app/cli/scaffold/templates/factory.py

@@ -1,7 +1,7 @@
 from ..constants import WORKFLOW_MOCK_TYPE, WORKFLOW_RECORD_TYPE, WORKFLOW_TEST_TYPE
 from ..docker.workflow.builder import WorkflowBuilder
 from .constants import (
-  CUSTOM_CONFIGURE, CUSTOM_INIT, MAINTAINED_CONFIGURE, MAINTAINED_FIXTURES, MOCK_WORKFLOW_CUSTOM_FILES, MOCK_WORKFLOW_MAINTAINED_FILES, RECORD_WORKFLOW_CUSTOM_FILES, RECORD_WORKFLOW_MAINTAINED_FILES, TEST_WORKFLOW_CUSTOM_FILES, TEST_WORKFLOW_MAINTAINED_FILES
+  CUSTOM_CONFIGURE, CUSTOM_INIT, MAINTAINED_CONFIGURE, MAINTAINED_PUBLIC, MOCK_WORKFLOW_CUSTOM_FILES, MOCK_WORKFLOW_MAINTAINED_FILES, RECORD_WORKFLOW_CUSTOM_FILES, RECORD_WORKFLOW_MAINTAINED_FILES, TEST_WORKFLOW_CUSTOM_FILES, TEST_WORKFLOW_MAINTAINED_FILES
 )
 
 def custom_files(workflow: str, workflow_builder: WorkflowBuilder):
@@ -21,8 +21,8 @@ def custom_files(workflow: str, workflow_builder: WorkflowBuilder):
 
   # Fixtures are only relevant if the workflow is mock/test and if the service has a hostname
   if not workflow_builder.config.hostname:
-    if MAINTAINED_FIXTURES in files:
-      files.remove(MAINTAINED_FIXTURES)
+    if MAINTAINED_PUBLIC in files:
+      files.remove(MAINTAINED_PUBLIC)
 
   return files
 
@@ -40,7 +40,7 @@ def maintained_files(workflow: str, workflow_builder: WorkflowBuilder):
       files.append(MAINTAINED_CONFIGURE)
 
   if not workflow_builder.config.hostname:
-    if MAINTAINED_FIXTURES in files:
-      files.remove(MAINTAINED_FIXTURES)
+    if MAINTAINED_PUBLIC in files:
+      files.remove(MAINTAINED_PUBLIC)
 
   return files

stoobly_agent/app/cli/scaffold/templates/workflow/mock/bin/configure

@@ -10,11 +10,4 @@ url="$scheme://$hostname"
 
 if [ "$scheme" = 'http' -a "$port" != '80' ] || [ "$scheme" = 'https' -a "$port" != '443' ]; then
   url="$url:$port"
-fi
-
-# Match Rules
-echo "Configuring match rules"
-stoobly-agent config match set \
-  --method GET --method POST --method OPTIONS --method PUT --method DELETE \
-  --mode mock \
-  --pattern ".*?"
+fi

stoobly_agent/app/cli/scaffold/templates/workflow/mock/fixtures.yml

@@ -1,5 +1,5 @@
 # The following example matches requests for GET /users/d+ (e.g. /users/1) with the contents of user-1.json
-# Assumes that 'user-1.json' is created in the 'fixtures' folder
+# Assumes that 'user-1.json' is created in a 'fixtures' folder in the same directory as this file
 #GET:
 #  /users/d+:
 #    path: ./fixtures/user-1.json

stoobly_agent/app/cli/scaffold/templates/workflow/test/bin/configure

@@ -10,11 +10,4 @@ url="$scheme://$hostname"
 
 if [ "$scheme" = 'http' -a "$port" != '80' ] || [ "$scheme" = 'https' -a "$port" != '443' ]; then
   url="$url:$port"
-fi
-
-# Match
-echo "Configuring match rules"
-stoobly-agent config match set \
-  --method GET --method POST --method OPTIONS --method PUT --method DELETE \
-  --mode mock \
-  --pattern ".*?"
+fi

stoobly_agent/app/cli/scaffold/templates/workflow/test/fixtures.yml

@@ -1,5 +1,5 @@
 # The following example matches requests for GET /users/d+ (e.g. /users/1) with the contents of user-1.json
-# Assumes that 'user-1.json' is created in the 'fixtures' folder
+# Assumes that 'user-1.json' is created in a 'fixtures' folder in the same directory as this file
 #GET:
 #  /users/d+:
 #    path: ./fixtures/user-1.json

stoobly_agent/app/cli/scaffold/workflow_command.py

@@ -6,7 +6,7 @@ from stoobly_agent.lib.logger import Logger
 
 from .app import App
 from .config import Config
-from .constants import BIN_FOLDER_NAME, COMPOSE_TEMPLATE, CONFIG_FILE, ENV_FILE, FIXTURES_FOLDER_NAME
+from .constants import BIN_FOLDER_NAME, COMPOSE_TEMPLATE, CONFIG_FILE, ENV_FILE, PUBLIC_FOLDER_NAME
 from .docker.constants import DOCKER_COMPOSE_CUSTOM
 from .service_command import ServiceCommand
 
@@ -88,8 +88,8 @@ class WorkflowCommand(ServiceCommand):
     return services
 
   @property
-  def fixtures_dir_path(self):
-    return os.path.join(self.workflow_path, FIXTURES_FOLDER_NAME)
+  def public_dir_path(self):
+    return os.path.join(self.workflow_path, PUBLIC_FOLDER_NAME)
 
   @property
   def workflow_config(self):

stoobly_agent/app/cli/scaffold/workflow_create_command.py

@@ -67,11 +67,11 @@ class WorkflowCreateCommand(WorkflowCommand):
       return
 
     # Maintained files are files that will always be overwritten
-    maintained_workflow_files = maintained_files(self.workflow_name, workflow_builder)
+    maintained_workflow_files = maintained_files(template or self.workflow_name, workflow_builder)
     self.copy_files(templates_path, maintained_workflow_files, self.workflow_path)
 
     # Custom files are files that may be modified by the user
-    custom_workflow_files = custom_files(self.workflow_name, workflow_builder)
+    custom_workflow_files = custom_files(template or self.workflow_name, workflow_builder)
     self.copy_files_no_replace(templates_path, custom_workflow_files, self.workflow_path)
 
   def __write_docker_compose_file(self, **kwargs: BuildOptions):

stoobly_agent/app/cli/scaffold_cli.py

@@ -140,19 +140,19 @@ def mkcert(**kwargs):
   help="Scaffold a service",
 )
 @click.option('--app-dir-path', default=current_working_dir, help='Path to application directory.')
-@click.option('--detached', is_flag=True)
+@click.option('--detached', is_flag=True, help='Use isolated and non-persistent context directory.')
 @click.option('--env', multiple=True, help='Specify an environment variable.')
 @click.option('--force', is_flag=True, help='Overwrite maintained scaffolded service files.')
-@click.option('--hostname')
-@click.option('--port')
-@click.option('--priority', default='5.0', help='Determines the service run order.')
+@click.option('--hostname', help='Service hostname.')
+@click.option('--port', help='Service port.')
+@click.option('--priority', default=5, type=click.FloatRange(1.0, 9.0), help='Determines the service run order. Lower values run first.')
 @click.option('--proxy-mode', help='''
   Proxy mode can be "regular", "transparent", "socks5",
   "reverse:SPEC", or "upstream:SPEC". For reverse and
   upstream proxy modes, SPEC is host specification in
   the form of "http[s]://host[:port]".
 ''')
-@click.option('--scheme', type=click.Choice(['http', 'https']))
+@click.option('--scheme', type=click.Choice(['http', 'https']), help='Defaults to https if hostname is set.')
 @click.option('--workflow', multiple=True, type=click.Choice([WORKFLOW_MOCK_TYPE, WORKFLOW_RECORD_TYPE, WORKFLOW_TEST_TYPE]), help='Include pre-defined workflows.')
 @click.argument('service_name')
 def create(**kwargs):

stoobly_agent/app/models/factories/resource/local_db/helpers/tiebreak_scenario_request.py

@@ -20,7 +20,7 @@ def generate_session_id(query: dict):
 
   return hashlib.md5(b'.'.join(toks)).hexdigest()
 
-def reset():
+def reset_sessions():
   Cache.instance().clear(f".+\.{PREFIX}")
 
 def tiebreak_scenario_request(session_id: str, requests: List[Request]):

stoobly_agent/app/models/factories/resource/local_db/request_adapter.py

@@ -10,7 +10,7 @@ from stoobly_agent.app.models.types import RequestCreateParams, RequestDestroyPa
 from stoobly_agent.app.proxy.mock.custom_not_found_response_builder import CustomNotFoundResponseBuilder
 from stoobly_agent.app.proxy.mock.ignored_components_response_builder import IgnoreComponentsResponseBuilder
 from stoobly_agent.app.proxy.record.joined_request import JoinedRequest
-from stoobly_agent.config.constants import custom_headers
+from stoobly_agent.config.constants import custom_headers, query_params as request_query_params
 from stoobly_agent.lib.orm import ORM
 from stoobly_agent.lib.orm.request import Request
 from stoobly_agent.lib.orm.response import Response
@@ -73,7 +73,6 @@ class LocalDBRequestAdapter(LocalDBAdapter):
   def response(self, **query_params: RequestColumns) -> requests.Response:
     self.__adapt_scenario_id(query_params)
 
-    endpoint_promise = query_params.get('endpoint_promise')
     request = None
 
     if not query_params.get('request_id'):
@@ -85,16 +84,19 @@ class LocalDBRequestAdapter(LocalDBAdapter):
       requests = self.__request_orm.where_for(**request_columns).get()
 
       if 'scenario_id' in query_params:
-        # TODO: Would need an additional ID to distinguish different scenario sessions
-        session_id = generate_session_id(request_columns) 
-
         if len(requests) > 1:
-          request = tiebreak_scenario_request(session_id, requests)
+          session_id = query_params.get(request_query_params.SESSION_ID)
+          request_session_id_components = { **request_columns }
+
+          if session_id:
+            request_session_id_components[request_query_params.SESSION_ID] = session_id
+
+          # When multiple requests are matched for a scenario, return them in sequence
+          request_session_id = generate_session_id(request_session_id_components)
+          request = tiebreak_scenario_request(request_session_id, requests)
+          access_request(request_session_id, request.id)
         else:
           request = requests.last()
-
-        if request:
-          access_request(session_id, request.id)
       else:
         request = requests.last()
     else:
@@ -104,6 +106,7 @@ class LocalDBRequestAdapter(LocalDBAdapter):
         request = None
 
     if not request:
+      endpoint_promise = query_params.get(request_query_params.ENDPOINT_PROMISE)
       return self.__handle_request_not_found(endpoint_promise) 
 
     response_record = request.response
@@ -315,8 +318,8 @@ class LocalDBRequestAdapter(LocalDBAdapter):
     return candidates.get()
 
   def __filter_request_response_columns(self, request_columns: RequestCreateParams):
-    if request_columns.get('endpoint_promise'):
-      del request_columns['endpoint_promise']
+    if request_columns.get(request_query_params.ENDPOINT_PROMISE):
+      del request_columns[request_query_params.ENDPOINT_PROMISE]
 
     if request_columns.get('infer'):
       del request_columns['infer']
@@ -327,6 +330,9 @@ class LocalDBRequestAdapter(LocalDBAdapter):
     if request_columns.get('retry'):
       del request_columns['retry']
 
+    if request_columns.get(request_query_params.SESSION_ID):
+      del request_columns[request_query_params.SESSION_ID]
+
   def __request(self, request_id: str):
     if self.validate_uuid(request_id):
       return self.__request_orm.find_by(uuid=request_id)

stoobly_agent/app/models/types/request.py

@@ -29,5 +29,4 @@ class RequestShowParams(TypedDict):
   headers: bool
   project_id: str
   query_params: bool
-  response: bool
-
+  response: bool

stoobly_agent/app/proxy/context.py

@@ -12,6 +12,10 @@ class InterceptContext():
   def flow(self):
     return self.__flow
 
+  @flow.setter
+  def flow(self, v: MitmproxyHTTPFlow):
+    self.__flow = v
+
   @property
   def intercept_settings(self):
     return self.__intercept_settings

stoobly_agent/app/proxy/handle_mock_service.py

@@ -3,7 +3,7 @@ import pdb
 import requests
 import time
 
-from mitmproxy.http import Request as MitmproxyRequest
+from mitmproxy.http import HTTPFlow as MitmproxyHTTPFlow, Request as MitmproxyRequest
 from typing import Callable, TypedDict
 
 from stoobly_agent.app.models.request_model import RequestModel
@@ -18,7 +18,8 @@ from .mock.eval_fixtures_service import eval_fixtures
 from .mock.eval_request_service import inject_eval_request
 from .utils.allowed_request_service import get_active_mode_policy
 from .utils.request_handler import reverse_proxy
-from .utils.response_handler import bad_request, pass_on
+from .utils.response_handler import bad_request, enable_cors, pass_on
+from .utils.rewrite import rewrite_request, rewrite_response
 
 LOG_ID = 'Mock'
 
@@ -26,28 +27,51 @@ class MockOptions(TypedDict):
     failure: Callable
     ignored_components: list 
     infer: bool
+    no_rewrite: bool
     success: Callable
 
+def handle_request_mock_generic_without_rewrite(context: MockContext, **options: MockOptions):
+    options['no_rewrite'] = True
+    handle_request_mock_generic(context, **options)
+
 ###
 #
+# 1. Rewrites mock request by default
+# 2. BEFORE_MOCK gets triggered
+# 3. AFTER_MOCK gets triggered
+#
 # @param request [mitmproxy.http.Request]
 # @param settings [Dict]
 #
 def handle_request_mock_generic(context: MockContext, **options: MockOptions):
-    __mock_hook(lifecycle_hooks.BEFORE_MOCK, context)
-
+    handle_error = options['error'] if 'error' in options and callable(options['error']) else None
+    handle_failure = options['failure'] if 'failure' in options and callable(options['failure']) else None
+    handle_success = options['success'] if 'success' in options and callable(options['success']) else None
     intercept_settings = context.intercept_settings
     request: MitmproxyRequest = context.flow.request
-    request_model = RequestModel(intercept_settings.settings)
+    res = None
 
     policy = get_active_mode_policy(request, intercept_settings)
+    if policy == mock_policy.NONE:
+        if handle_error:
+            res = handle_error(context)
 
-    rewrite_rules = intercept_settings.mock_rewrite_rules
-    if len(rewrite_rules) > 0:
-        # Rewrite request with paramter rules for mock
-        request: MitmproxyRequest = context.flow.request
-        request_facade = MitmproxyRequestFacade(request)
-        request_facade.with_parameter_rules(rewrite_rules).with_url_rules(rewrite_rules).rewrite()
+        return pass_on(context.flow, res)
+
+    if policy not in [mock_policy.ALL, mock_policy.FOUND]:
+        if handle_error:
+            res = handle_error(context)
+
+        return bad_request(
+            context.flow,
+            "Valid env MOCK_POLICY: %s, Got: %s" %
+            ([mock_policy.ALL, mock_policy.FOUND, mock_policy.NONE], policy)
+        )
+
+    if not options.get('no_rewrite'):
+        __rewrite_request(context)
+
+    __mock_hook(lifecycle_hooks.BEFORE_MOCK, context)
 
     # If ignore rules are set, then ignore specified request parameters
     ignore_rules = intercept_settings.ignore_rules
@@ -55,48 +79,33 @@ def handle_request_mock_generic(context: MockContext, **options: MockOptions):
         request_facade = MitmproxyRequestFacade(request)
         _ignore_rules = request_facade.select_parameter_rules(ignore_rules)
         ignored_components = rewrite_rules_to_ignored_components(_ignore_rules)
-        options['ignored_components'] += ignored_components  if 'ignored_components' in options else ignored_components
+        options['ignored_components'] += ignored_components if 'ignored_components' in options else ignored_components 
 
-    handle_success = options['success'] if 'success' in options and callable(options['success']) else None
-    handle_failure = options['failure'] if 'failure' in options and callable(options['failure']) else None
-    
+    request_model = RequestModel(intercept_settings.settings)
     eval_request = inject_eval_request(request_model, intercept_settings)
- 
-    if policy == mock_policy.NONE:
-        if handle_failure:
-            res = handle_failure(context)
-    elif policy == mock_policy.ALL:
+
+    if policy == mock_policy.ALL:
         res = eval_request_with_retry(context, eval_request, **options) 
 
         context.with_response(res)
-
-        if handle_success:
-            # TODO: rewrite response, see #332
-            res = handle_success(context) or res
     elif policy == mock_policy.FOUND:
         res = eval_request_with_retry(context, eval_request, **options) 
 
         context.with_response(res)
 
         if res.status_code in [custom_response_codes.NOT_FOUND, custom_response_codes.IGNORE_COMPONENTS]:
-            if handle_failure:
-                try:
-                    res = handle_failure(context)
-                except RuntimeError:
-                    # Do nothing, return custom error response
-                    pass
-        else:
-            if handle_success:
-                # TODO: rewrite response, see #332
-                res = handle_success(context) or res
+            try:
+                res = __handle_found_policy(context)
+            except RuntimeError:
+                # Do nothing, return custom error response
+                pass
+        
+    if res.status_code == custom_response_codes.NOT_FOUND:
+        if handle_failure:
+            res = handle_failure(context) or res
     else:
-        return bad_request(
-            context.flow,
-            "Valid env MOCK_POLICY: %s, Got: %s" %
-            ([mock_policy.ALL, mock_policy.FOUND, mock_policy.NONE], policy)
-        )
-
-    __mock_hook(lifecycle_hooks.AFTER_MOCK, context)
+        if handle_success:
+            res = handle_success(context) or res
 
     return pass_on(context.flow, res)
 
@@ -130,6 +139,11 @@ def handle_request_mock(context: MockContext):
         success=lambda context: __handle_mock_success(context)
     )
 
+###
+#
+# 1. Rewrites mock response (if mock found)
+# 2. AFTER_MOCK gets triggered (if mock found)
+#
 def handle_response_mock(context: MockContext):
     response = context.flow.response
     request_key = response.headers.get(custom_headers.MOCK_REQUEST_KEY)
@@ -138,13 +152,22 @@ def handle_response_mock(context: MockContext):
         request = context.flow.request
         Logger.instance(LOG_ID).info(f"{bcolors.OKCYAN}Mocked{bcolors.ENDC} {request.url} -> {request_key}")
 
-def __handle_mock_success(context: MockContext) -> None:
-    if os.environ.get(env_vars.AGENT_SIMULATE_LATENCY):
-        response = context.response
-        start_time = context.start_time
-        __simulate_latency(response.headers.get(custom_headers.RESPONSE_LATENCY), start_time)
+        __rewrite_response(context)
+        __mock_hook(lifecycle_hooks.AFTER_MOCK, context)
 
 def __handle_mock_failure(context: MockContext) -> None:
+    flow = context.flow
+    request = flow.request
+
+    if request.method.upper() != 'OPTIONS':
+        return False
+        
+    # Default OPTIONS request to allow CORS
+    enable_cors(flow)
+
+    return True
+
+def __handle_found_policy(context: MockContext) -> None:
     req = context.flow.request
     intercept_settings = context.intercept_settings
     upstream_url = intercept_settings.upstream_url
@@ -160,6 +183,30 @@ def __handle_mock_failure(context: MockContext) -> None:
 
     reverse_proxy(req, upstream_url, {})
 
+def __handle_mock_success(context: MockContext) -> None:
+    if os.environ.get(env_vars.AGENT_SIMULATE_LATENCY):
+        response = context.response
+        start_time = context.start_time
+        __simulate_latency(response.headers.get(custom_headers.RESPONSE_LATENCY), start_time)
+
+def __rewrite_request(context: MockContext):
+    # Rewrite request with paramter rules for mock
+
+    intercept_settings = context.intercept_settings
+    rewrite_rules = intercept_settings.mock_rewrite_rules
+
+    if len(rewrite_rules) > 0:
+        rewrite_request(context.flow, rewrite_rules) 
+
+def __rewrite_response(context: MockContext):
+    # Rewrite request with paramter rules for mock
+
+    intercept_settings = context.intercept_settings
+    rewrite_rules = intercept_settings.mock_rewrite_rules
+
+    if len(rewrite_rules) > 0:
+        rewrite_response(context.flow, rewrite_rules) 
+
 ###
 #
 # Try to simulate expected response latency