stoobly-agent 1.8.5__tar.gz → 1.9.1__tar.gz

{stoobly_agent-1.8.5 → stoobly_agent-1.9.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: stoobly-agent
-Version: 1.8.5
+Version: 1.9.1
 Summary: Record, mock, and test HTTP(s) requests. CLI agent for Stoobly
 License: Apache-2.0
 Author: Matt Le

{stoobly_agent-1.8.5 → stoobly_agent-1.9.1}/pyproject.toml

@@ -10,7 +10,7 @@ description = "Record, mock, and test HTTP(s) requests. CLI agent for Stoobly"
 license = "Apache License 2.0"
 name = "stoobly-agent"
 readme = "README.md"
-version = "1.8.5"
+version = "1.9.1"
 
 [tool.poetry.dependencies]
 python = "^3.10"

{stoobly_agent-1.8.5 → stoobly_agent-1.9.1}/stoobly_agent/__init__.py

@@ -1,2 +1,2 @@
 COMMAND = 'stoobly-agent'
-VERSION = '1.8.5'
+VERSION = '1.9.1'

{stoobly_agent-1.8.5 → stoobly_agent-1.9.1}/stoobly_agent/app/api/configs_controller.py

@@ -4,7 +4,7 @@ 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_intercept_active_update, handle_policy_update, handle_project_update, handle_scenario_update
+    context as handle_context, handle_intercept_active_update, handle_order_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
@@ -41,7 +41,7 @@ class ConfigsController:
             )
         elif active_mode == mode.RECORD:
             context.render(
-                json = [record_policy.ALL, record_policy.FOUND, record_policy.NOT_FOUND, record_policy.OVERWRITE],
+                json = [record_policy.ALL, record_policy.API, record_policy.FOUND, record_policy.NOT_FOUND],
                 status = 200
             )
         elif active_mode == mode.REPLAY:
@@ -112,7 +112,7 @@ class ConfigsController:
         _handle_context = handle_context()
 
         handle_intercept_active_update(settings, _handle_context)
-        handle_policy_update(settings, _handle_context)
+        handle_order_update(settings, _handle_context)
         handle_project_update(settings, _handle_context)
         handle_scenario_update(settings, _handle_context)
 

{stoobly_agent-1.8.5 → stoobly_agent-1.9.1}/stoobly_agent/app/cli/helpers/handle_config_update_service.py

@@ -6,7 +6,7 @@ from stoobly_agent.app.models.scenario_model import ScenarioModel
 from stoobly_agent.app.settings.constants import intercept_mode
 from stoobly_agent.app.settings.constants.intercept_mode import Mode
 from stoobly_agent.app.settings import ProxySettings, Settings
-from stoobly_agent.config.constants import record_policy
+from stoobly_agent.config.constants import record_order
 from stoobly_agent.lib.api.keys.project_key import ProjectKey
 from stoobly_agent.lib.api.keys.scenario_key import ScenarioKey
 
@@ -28,15 +28,15 @@ def handle_intercept_active_update(new_settings: Settings, context: Context = No
 
   if not was_active and is_active:
     if _mode == intercept_mode.RECORD:
-      new_policy = data_rule.record_policy
+      new_order = data_rule.record_order
       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,
+        if new_order == record_order.OVERWRITE:
+          # If order is overwrite when recording, whenever intercept is enabled,
           # set active scenario to be overwritable
           scenario_model.update(_scenario_key.id, **{ 'overwritable': True })[1]
         else:
@@ -44,16 +44,16 @@ def handle_intercept_active_update(new_settings: Settings, context: Context = No
   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
+      old_order = __data_rule(old_proxy_settings).record_order
 
-      if old_policy == record_policy.OVERWRITE:
+      if old_order == record_order.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,
+          # If order is overwrite when recording, whenever intercept is disabled,
           # set active scenario to not be overwritable
           scenario_model.update(_scenario_key.id, **{ 'overwritable': False })[1]
     elif _mode == intercept_mode.MOCK:
@@ -73,7 +73,7 @@ def handle_scenario_update(new_settings: Settings, context = None):
   if old_scenario_key != new_scenario_key:
     data_rule = __data_rule(new_settings.proxy)
 
-    if data_rule.record_policy == record_policy.OVERWRITE:
+    if data_rule.record_order == record_order.OVERWRITE:
       scenario_model = ScenarioModel(new_settings)
 
       _old_scenario_key = __parse_scenario_key(old_scenario_key)
@@ -108,7 +108,7 @@ def handle_project_update(new_settings: Settings, context: Context = None):
     if _project_key.id != _new_scenario_key.project_id:
       data_rule.scenario_key = None
 
-def handle_policy_update(new_settings: Settings, context: Context = None):
+def handle_order_update(new_settings: Settings, context: Context = None):
   data_rule = __data_rule(new_settings.proxy)
   _mode = context['mode'] if context and context.get('mode') else new_settings.proxy.intercept.mode
 
@@ -116,10 +116,10 @@ def handle_policy_update(new_settings: Settings, context: Context = None):
   old_data_rule = __data_rule(old_proxy_settings)
 
   if _mode == intercept_mode.RECORD:
-    new_policy = data_rule.record_policy
-    old_policy = old_data_rule.record_policy
+    new_order = data_rule.record_order
+    old_order = old_data_rule.record_order
 
-    if new_policy != old_policy and old_policy == record_policy.OVERWRITE:
+    if new_order != old_order and old_order == record_order.OVERWRITE:
       scenario_key = data_rule.scenario_key
       _scenario_key = __parse_scenario_key(scenario_key)
 

{stoobly_agent-1.8.5 → stoobly_agent-1.9.1}/stoobly_agent/app/cli/intercept_cli.py

@@ -3,10 +3,10 @@ import pdb
 import sys
 
 from stoobly_agent.app.settings import Settings
-from stoobly_agent.config.constants import mode, mock_policy, record_policy, replay_policy
+from stoobly_agent.config.constants import mode, mock_policy, record_order, record_policy, replay_policy
 from stoobly_agent.lib.api.keys.project_key import ProjectKey
 
-from .helpers.handle_config_update_service import handle_intercept_active_update, handle_policy_update
+from .helpers.handle_config_update_service import handle_intercept_active_update, handle_order_update
 
 settings = Settings.instance()
 
@@ -17,16 +17,25 @@ if settings.cli.features.remote:
 
 active_mode = settings.proxy.intercept.mode
 
+def __get_order_options(active_mode):
+    if active_mode == mode.RECORD:
+        return [record_order.APPEND, record_order.OVERWRITE]
+    else:
+        return []
+
 def __get_policy_options(active_mode):
     if active_mode == mode.MOCK:
         return [mock_policy.ALL, mock_policy.FOUND]
     elif active_mode == mode.RECORD:
-        return [record_policy.ALL, record_policy.FOUND, record_policy.NOT_FOUND,record_policy.OVERWRITE]
+        return [record_policy.ALL, record_policy.API, record_policy.FOUND, record_policy.NOT_FOUND]
     elif active_mode == mode.REPLAY:
         return [replay_policy.ALL]
     elif active_mode == mode.TEST:
         return [mock_policy.FOUND]
+    else:
+        return []
 
+order_options = __get_order_options(active_mode)
 policy_options = __get_policy_options(active_mode)
 
 @click.group(
@@ -69,11 +78,12 @@ def disable(**kwargs):
     help="Configure intercept"
 )
 @click.option('--mode', type=click.Choice(mode_options))
+@click.option('--order', type=click.Choice(order_options))
 @click.option('--policy', type=click.Choice(policy_options))
 def configure(**kwargs):
-    settings = Settings.instance()
+    settings: Settings = Settings.instance()
 
-    if not kwargs['mode'] and not kwargs['policy']:
+    if not kwargs['mode'] and not kwargs['order'] and not kwargs['policy']:
         print("Error: Missing an option")
         sys.exit(1)
 
@@ -89,6 +99,21 @@ def configure(**kwargs):
 
     _mode = kwargs['mode'] or settings.proxy.intercept.mode 
 
+    if kwargs['order']:
+        active_mode = settings.proxy.intercept.mode
+
+        if active_mode == mode.RECORD:
+            project_key = ProjectKey(settings.proxy.intercept.project_key)
+            data_rule = settings.proxy.data.data_rules(project_key.id)
+            data_rule.record_order = kwargs['order']
+        else:
+            print("Error: set --mode to a intercept mode that supports the order option", file=sys.stderr)
+            sys.exit(1)
+
+        handle_order_update(settings)
+
+        print(f"Updating {_mode} order to {kwargs['order']}")
+
     if kwargs['policy']:
         active_mode = settings.proxy.intercept.mode
         valid_policies = __get_policy_options(active_mode)
@@ -109,8 +134,6 @@ 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']}")
 
     settings.commit()

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

@@ -0,0 +1,6 @@
+FROM stoobly/agent:1.9
+
+ARG USER_ID
+
+# Change user id of stoobly user to that of host's user id
+RUN if [ -n "$USER_ID" ]; then usermod -u $USER_ID stoobly; fi

{stoobly_agent-1.8.5 → stoobly_agent-1.9.1}/stoobly_agent/app/cli/scaffold/templates/app/.Makefile

@@ -143,6 +143,12 @@ scenario/list:
 	@export EXEC_COMMAND=.list && \
 	export EXEC_OPTIONS="$(options)" && \
 	$(stoobly_exec)
+scenario/overwrite:
+# Overwrite a scenario
+	@export EXEC_COMMAND=.overwrite && \
+	export EXEC_OPTIONS="$(options)" && \
+	export EXEC_ARGS="$(key)" && \
+	$(stoobly_exec)
 scenario/reset:
 # Resets a scenario to its last snapshot
 	@export EXEC_COMMAND=.reset && \:

stoobly_agent-1.9.1/stoobly_agent/app/cli/scaffold/templates/build/workflows/exec/.overwrite

@@ -0,0 +1,13 @@
+#!/bin/bash
+
+scenario_key=$1
+
+if [ ! "$scenario_key" ]; then 
+    echo "Missing argument 'scenario_key'"
+    stoobly-agent scenario list
+    exit
+fi
+
+stoobly-agent config scenario set "$scenario_key"
+stoobly-agent intercept configure --mode record --order overwrite
+stoobly-agent intercept disable

{stoobly_agent-1.8.5 → stoobly_agent-1.9.1}/stoobly_agent/app/cli/scaffold/templates/workflow/record/bin/configure

@@ -22,6 +22,14 @@ stoobly-agent config firewall set \
 
 # Rewrite
 echo "Configuring rewrite rules..."
+stoobly-agent config rewrite set \
+  --method GET --method POST --method OPTIONS --method PUT --method DELETE \
+  --mode record \
+  --name authorization \
+  --pattern "$origin/?.*?" \
+  --type Header \
+  --value '' \
+
 stoobly-agent config rewrite set \
   --method GET --method POST --method OPTIONS --method PUT --method DELETE \
   --mode record \
@@ -29,3 +37,11 @@ stoobly-agent config rewrite set \
   --pattern "$origin/?.*?" \
   --type Header \
   --value '' \
+
+stoobly-agent config rewrite set \
+  --method GET --method POST --method OPTIONS --method PUT --method DELETE \
+  --mode record \
+  --name set-cookie \
+  --pattern "$origin/?.*?" \
+  --type 'Response Header' \
+  --value '' \

{stoobly_agent-1.8.5 → stoobly_agent-1.9.1}/stoobly_agent/app/models/adapters/joined_request_adapter.py

@@ -46,7 +46,7 @@ class JoinedRequestAdapter():
   def build_request_string(self):
     request_string = RequestString(None)
 
-    delimitter = RequestStringCLRF.encode()
+    delimitter = RequestStringCLRF
     request_string_toks = self.__split_joined_request_string[0].split(delimitter)
     request_string.set(self.raw_request_string or delimitter.join(request_string_toks[1:]))
     request_string.control = request_string_toks[0]

{stoobly_agent-1.8.5 → stoobly_agent-1.9.1}/stoobly_agent/app/models/factories/resource/local_db/helpers/log.py

@@ -8,6 +8,7 @@ from stoobly_agent.config.data_dir import DataDir
 from stoobly_agent.lib.logger import bcolors, Logger
 
 from .log_event import LogEvent
+from .request_snapshot import RequestSnapshot
 from .snapshot_types import DELETE_ACTION, PUT_ACTION, Resource
 
 EVENT_DELIMITTER = "\n"
@@ -67,6 +68,25 @@ class Log():
     contents = self.read()
     return self.build_raw_events(contents) 
 
+  @property
+  def scenario_inverted_index(self):
+    index = {}
+
+    def handle_snapshot(snapshot: RequestSnapshot):
+      request_uuid = snapshot.uuid
+      if not request_uuid in index:
+        index[request_uuid] = []
+
+      index[request_uuid].append(event.resource_uuid)
+
+    for event in self.target_events:
+      if not event.is_scenario():
+        continue 
+
+      event.snapshot().iter_request_snapshots(handle_snapshot)
+
+    return index
+
   @property
   def unprocessed_events(self) -> List[LogEvent]:
     events = self.events
@@ -197,12 +217,18 @@ class Log():
 
       resource_index[event.resource_uuid].append(event)
 
+    scenario_inverted_index = self.scenario_inverted_index
+
     pruned_events = self.collapse(events)
     for event in pruned_events:
       snapshot = event.snapshot()
       snapshot_exists = snapshot.exists
 
       if event.action == DELETE_ACTION or not snapshot_exists:
+        if event.is_request() and event.resource_uuid in scenario_inverted_index:
+          # If a request is deleted, only prune if it's not also a part of a scenario
+          continue
+
         Logger.instance(LOG_ID).info(f"{bcolors.OKBLUE}Removing{bcolors.ENDC} {event.resource} {event.resource_uuid}")
 
         resource_events: List[LogEvent] = resource_index[event.resource_uuid]
@@ -218,8 +244,16 @@ class Log():
           removed_events[event.uuid] = True
 
         if event.action == DELETE_ACTION and snapshot_exists: 
-          if not dry_run:
+          if dry_run:
+            continue
+
+          if event.is_scenario():
+            # We still need to check each request in a scenario to make sure another scenario does not depend on it
+            snapshot.remove(lambda snapshot: self.remove_request_snapshot(snapshot, scenario_inverted_index))
+          elif event.is_request():
+            # We have already checked that a scenario does not depend on the request above
             snapshot.remove()
+
           Logger.instance(LOG_ID).info(f"Removing {event.resource} snapshot")
 
   def build_raw_events(self, contents: str) -> List[str]:
@@ -242,6 +276,8 @@ class Log():
     '''
     Remove DELETE events where the last processed event was a PUT
     '''
+
+    # Build an index such that if the last event is DELETE_ACTION, then it will NOT exist in the index
     index = {}
     for event in processed_events:
       if event.action == PUT_ACTION:
@@ -249,10 +285,18 @@ class Log():
       elif event.action == DELETE_ACTION:
         if event.resource_uuid in index:
           del index[event.resource_uuid]
+
+    scenario_inverted_index = self.scenario_inverted_index
+
+    def keep(e: LogEvent):
+      if e.action != DELETE_ACTION:
+        return True
+      
+      return e.action == DELETE_ACTION and e.is_request() and e.resource_uuid in scenario_inverted_index
     
     return list(
       filter(
-        lambda e: e.action != DELETE_ACTION or (e.action == DELETE_ACTION and e.resource_uuid in index), 
+        keep,
         unprocessed_events
       )
     )
@@ -289,6 +333,12 @@ class Log():
     remaining_events = self.remove_dangling_events(processed_events, unprocessed_events)
     return list(filter(lambda e: e.uuid not in remaining_version_uuids, remaining_events))
 
+  def remove_request_snapshot(self, snapshot: RequestSnapshot, scenario_inverted_index: dict = None):
+    scenario_inverted_index = scenario_inverted_index or self.scenario_inverted_index
+
+    if snapshot.uuid not in scenario_inverted_index:
+      snapshot.remove()
+
   # Rotate log to history
   def rotate(self):
     if not os.path.exists(self.__log_file_path):

{stoobly_agent-1.8.5 → stoobly_agent-1.9.1}/stoobly_agent/app/models/factories/resource/local_db/helpers/request_snapshot.py

@@ -1,5 +1,4 @@
 import os
-import shutil
 
 from stoobly_agent.app.models.adapters.orm import JoinedRequestStringAdapter
 from stoobly_agent.lib.orm.request import Request

{stoobly_agent-1.8.5 → stoobly_agent-1.9.1}/stoobly_agent/app/models/factories/resource/local_db/helpers/scenario_snapshot.py

@@ -1,7 +1,6 @@
 import json
 import os
 import pdb
-import shutil
 
 from typing import Callable, List
 
@@ -128,12 +127,12 @@ class ScenarioSnapshot(Snapshot):
     if os.path.exists(metadata_file_path):
       os.remove(metadata_file_path)
 
-  def remove_requests(self):
+  def remove_requests(self, handle_remove_request_snapshot = None):
     requests_file_path = self.requests_path
 
     if os.path.exists(requests_file_path):
       # A request only every belongs to one scenario
-      self.iter_request_snapshots(self.__handle_remove_requests)
+      self.iter_request_snapshots(handle_remove_request_snapshot or self.__handle_remove_requests)
 
       os.remove(requests_file_path)
 

{stoobly_agent-1.8.5 → stoobly_agent-1.9.1}/stoobly_agent/app/models/factories/resource/local_db/helpers/snapshot_service.py

@@ -14,15 +14,17 @@ def snapshot_request(request: Request, **options: RequestSnapshotOptions):
     return
 
   snapshot = RequestSnapshot(request.uuid)
-
   snapshot.backup()
 
+  log = Log()
+
   if action == PUT_ACTION:
     snapshot.write(request, **options)
   elif action == DELETE_ACTION:
-    snapshot.remove()
+    inverted_scenario_index = log.scenario_inverted_index
 
-  log = Log()
+    # If a scenario currently depends on this request, we can't remove the snapshot until the scenario is also removed
+    log.remove_request_snapshot(snapshot, inverted_scenario_index) 
 
   try:
     if action == PUT_ACTION:
@@ -40,7 +42,6 @@ def snapshot_scenario(scenario: Scenario, **options):
     return
 
   snapshot = ScenarioSnapshot(scenario.uuid)
-
   snapshot.backup_metadata()
 
   if action == PUT_ACTION:
@@ -50,13 +51,13 @@ def snapshot_scenario(scenario: Scenario, **options):
 
   snapshot.backup_requests()
 
+  log = Log()
+  inverted_scenario_index = log.scenario_inverted_index
   if action == PUT_ACTION:
-    snapshot.remove_requests()
+    snapshot.remove_requests(lambda snapshot: log.remove_request_snapshot(snapshot, inverted_scenario_index))
     snapshot.write_requests(scenario, **options)
   elif action == DELETE_ACTION:
-    snapshot.remove_requests()
-
-  log = Log()
+    snapshot.remove_requests(lambda snapshot: log.remove_request_snapshot(snapshot, inverted_scenario_index))
 
   if action == PUT_ACTION:
     log.put(scenario)

{stoobly_agent-1.8.5 → stoobly_agent-1.9.1}/stoobly_agent/app/models/helpers/apply.py

@@ -1,9 +1,11 @@
 import pdb
 
+from typing import List
+
 from stoobly_agent.app.models.factories.resource.local_db.helpers.log import Log
 from stoobly_agent.app.models.factories.resource.local_db.helpers.request_snapshot import RequestSnapshot
 from stoobly_agent.app.models.factories.resource.local_db.helpers.scenario_snapshot import ScenarioSnapshot
-from stoobly_agent.app.proxy.record import REQUEST_STRING_CLRF, RequestStringControl
+from stoobly_agent.app.proxy.record import REQUEST_STRING_CLRF
 from stoobly_agent.app.settings import Settings
 from stoobly_agent.lib.logger import bcolors
 
@@ -71,6 +73,7 @@ class Apply():
       if results:
         status = results[1]
         if status == 0 or status >= 400:
+          self.__logger(f"{bcolors.FAIL}Error{bcolors.ENDC} {results[0]}")
           completed = False
           break
 
@@ -218,14 +221,16 @@ class Apply():
 
     snapshot_requests = {}
 
-    request_snapshots = snapshot.request_snapshots
+    request_snapshots: List[RequestSnapshot] = snapshot.request_snapshots
     for request_snapshot in request_snapshots:
       raw_request = request_snapshot.request
-      
-      toks = raw_request.split(REQUEST_STRING_CLRF.encode(), 1)
 
+      if not raw_request:
+        return f"{request_snapshot.path} is missing", 400
+        
+      toks = raw_request.split(REQUEST_STRING_CLRF, 1)
       if len(toks) != 2:
-        return f"{snapshot.requests_path} contains an invalid request", 400
+        return f"{request_snapshot.path} contains an invalid request", 400
 
       uuid = request_snapshot.uuid
       res, status = self.__put_request(uuid, raw_request, scenario_id=scenario['id'])

{stoobly_agent-1.8.5 → stoobly_agent-1.9.1}/stoobly_agent/app/proxy/handle_mock_service.py

@@ -3,7 +3,7 @@ import pdb
 import requests
 import time
 
-from mitmproxy.http import HTTPFlow as MitmproxyHTTPFlow, Request as MitmproxyRequest
+from mitmproxy.http import Request as MitmproxyRequest
 from typing import Callable, TypedDict
 
 from stoobly_agent.app.models.request_model import RequestModel
@@ -64,7 +64,7 @@ def handle_request_mock_generic(context: MockContext, **options: MockOptions):
 
         return bad_request(
             context.flow,
-            "Valid env MOCK_POLICY: %s, Got: %s" %
+            "Valid mock policies: %s, Got: %s" %
             ([mock_policy.ALL, mock_policy.FOUND, mock_policy.NONE], policy)
         )
 
@@ -88,18 +88,20 @@ def handle_request_mock_generic(context: MockContext, **options: MockOptions):
         res = eval_request_with_retry(context, eval_request, **options) 
 
         context.with_response(res)
+        res = __after_mock_not_found(context)
     elif policy == mock_policy.FOUND:
         res = eval_request_with_retry(context, eval_request, **options) 
 
         context.with_response(res)
+        res = __after_mock_not_found(context)
 
-        if res.status_code in [custom_response_codes.NOT_FOUND, custom_response_codes.IGNORE_COMPONENTS]:
+        if res.status_code == custom_response_codes.NOT_FOUND:
             try:
                 return __handle_found_policy(context) # Continue proxying the request
             except RuntimeError:
                 # Do nothing, return custom error response
-                pass
-    
+                pass 
+
     if res.status_code == custom_response_codes.NOT_FOUND:
         if handle_failure:
             res = handle_failure(context) or res
@@ -241,4 +243,15 @@ def __mock_hook(hook: str, context: MockContext):
     lifecycle_hooks_module = intercept_settings.lifecycle_hooks
 
     if hook in lifecycle_hooks_module:
-        lifecycle_hooks_module[hook](context)
+        lifecycle_hooks_module[hook](context)
+
+def __after_mock_not_found(context: MockContext):
+    res = context.response
+
+    if res.status_code == custom_response_codes.NOT_FOUND:
+        __mock_hook(lifecycle_hooks.AFTER_MOCK_NOT_FOUND, context)
+
+        # context.response may have been modified by the hook
+        res = context.response
+
+    return res

{stoobly_agent-1.8.5 → stoobly_agent-1.9.1}/stoobly_agent/app/proxy/handle_record_service.py

@@ -9,7 +9,7 @@ from stoobly_agent.app.settings.constants.mode import TEST
 from stoobly_agent.app.models.request_model import RequestModel
 from stoobly_agent.app.proxy.intercept_settings import InterceptSettings
 from stoobly_agent.config.constants.env_vars import ENV
-from stoobly_agent.config.constants import lifecycle_hooks, record_policy
+from stoobly_agent.config.constants import lifecycle_hooks, record_order, record_policy
 from stoobly_agent.lib.logger import Logger
 
 from .constants import custom_response_codes
@@ -17,6 +17,7 @@ from .mock.eval_request_service import inject_eval_request
 from .record.context import RecordContext
 from .record.overwrite_scenario_service import overwrite_scenario
 from .record.upload_request_service import inject_upload_request
+from .replay.body_parser_service import is_json, is_xml
 from .utils.allowed_request_service import get_active_mode_policy
 from .utils.response_handler import bad_request, disable_transfer_encoding 
 from .utils.rewrite import rewrite_request_response
@@ -42,6 +43,13 @@ def handle_response_record(context: RecordContext):
 
     if active_record_policy == record_policy.ALL:
         __record_request(context, request_model)
+    elif active_record_policy == record_policy.API:
+        response = flow.response
+        content_type: str = response.headers.get('content-type')
+
+        if content_type:
+            if is_json(content_type) or is_xml(content_type) or content_type.startswith('text/plain'):
+                __record_request(context, request_model)
     elif active_record_policy == record_policy.FOUND:
         res = inject_eval_request(request_model, intercept_settings)(request, [])
 
@@ -52,16 +60,12 @@ def handle_response_record(context: RecordContext):
 
         if res.status_code == custom_response_codes.NOT_FOUND:
             __record_request(context, request_model)
-    elif active_record_policy == record_policy.OVERWRITE:
-        overwrite_scenario(intercept_settings.scenario_key)
-
-        __record_request(context, request_model)
     else:
         if active_record_policy != record_policy.NONE:
             return bad_request(
                 flow,
-                "Valid env RECORD_POLICY: %s, Got: %s" %
-                ([record_policy.ALL, record_policy.FOUND, record_policy.NOT_FOUND], active_record_policy)
+                "Valid record policies: %s, Got: %s" %
+                ([record_policy.ALL, record_policy.API, record_policy.FOUND, record_policy.NOT_FOUND], active_record_policy)
             )
 
 def __record_handler(context: RecordContext, request_model: RequestModel):
@@ -79,6 +83,11 @@ def __record_handler(context: RecordContext, request_model: RequestModel):
     context.flow = flow # Reset flow
 
 def __record_request(context: RecordContext, request_model: RequestModel):
+    intercept_settings = context.intercept_settings
+
+    if intercept_settings.order == record_order.OVERWRITE:
+        overwrite_scenario(intercept_settings.scenario_key)
+
     if os.environ.get(ENV) == TEST:
         __record_handler(context, request_model)
     else: