stoobly-agent 1.8.3__py3-none-any.whl → 1.8.5__py3-none-any.whl

stoobly_agent/__init__.py

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

stoobly_agent/app/api/headers_controller.py

@@ -86,7 +86,7 @@ class HeadersController:
 
         context.render(
             plain = '',
-            status = 200
+            status = 204
         )
 
     def __header_model(self, context: SimpleHTTPRequestHandler):

stoobly_agent/app/api/query_params_controller.py

@@ -103,13 +103,13 @@ class QueryParamsController:
 
         if context.filter_response(query_param, status):
             return
-        
+
         context.render(
-            plain = query_param,
-            status = status
+            plain = '',
+            status = 204
         )
 
     def __query_param_model(self, context: SimpleHTTPRequestHandler):
         access_token = context.headers.get('access-token')
         query_param_model = QueryParamModel(Settings.instance(), access_token=access_token)
-        return query_param_model
+        return query_param_model

stoobly_agent/app/cli/snapshot_cli.py

@@ -20,6 +20,7 @@ from stoobly_agent.lib.api.keys import RequestKey, ScenarioKey
 
 from .helpers.print_service import FORMATS, print_snapshots, select_print_options
 from .helpers.verify_raw_request_service import verify_raw_request
+from .types.snapshot_migration import SnapshotMigration
 
 @click.group(
     epilog="Run 'stoobly-agent project COMMAND --help' for more information on a command.",
@@ -91,6 +92,62 @@ def _list(**kwargs):
     if len(formatted_events):
       print_snapshots(formatted_events, **print_options)
 
+@snapshot.command(
+  help="Migrate snapshots."
+)
+@click.option('--scenario-key', help='Apply migration to specific scenario.')
+@click.argument('lifecycle_hooks_path')
+def migrate(**kwargs):
+  from runpy import run_path
+  from stoobly_agent.config.constants.lifecycle_hooks import BEFORE_MIGRATE
+
+  scenario_key = None
+  if kwargs['scenario_key']:
+    scenario_key = ScenarioKey(kwargs['scenario_key'])
+
+  try:
+    lifecycle_hooks = run_path(kwargs['lifecycle_hooks_path'])
+  except Exception as e:
+    lifecycle_hooks = {}
+
+  before_migrate_hook = None
+  if BEFORE_MIGRATE not in lifecycle_hooks:
+    sys.exit(1)
+  else:
+    before_migrate_hook = lifecycle_hooks[BEFORE_MIGRATE]
+
+  log = Log()
+
+  request_snapshots = []
+  for event in log.resource_events:
+    if event.is_scenario():
+      # If scenario_key is set, only consider that scenario
+      if scenario_key and event.resource_uuid != scenario_key.id:
+        continue
+
+      scenario_snapshot: ScenarioSnapshot = event.snapshot()
+      _request_snapshots = scenario_snapshot.request_snapshots
+
+      for request_snapshot in _request_snapshots:
+        # Scenario requests will have the same log event
+        request_snapshots.append((request_snapshot, event))
+    elif event.is_request():
+      # If scenario_key is set, only consider that scenario
+      if scenario_key:
+        continue
+
+      request_snapshot: RequestSnapshot = event.snapshot()
+      request_snapshots.append((request_snapshot, event))
+
+  snapshot_migrations = {}
+  for request_snapshot, event in request_snapshots:
+    if request_snapshot.uuid not in snapshot_migrations:
+      snapshot_migration = SnapshotMigration(request_snapshot, event)
+      snapshot_migrations[request_snapshot.uuid] = snapshot_migration
+      
+      if before_migrate_hook(snapshot_migration):
+        break
+
 @snapshot.command(
   help="Prune deleted snapshots."
 )

stoobly_agent/app/cli/types/snapshot_migration.py

@@ -0,0 +1,61 @@
+from mitmproxy.http import Request, Response
+
+from stoobly_agent.app.models.factories.resource.local_db.helpers.log import Log
+from stoobly_agent.app.models.factories.resource.local_db.helpers.log_event import LogEvent
+from stoobly_agent.app.models.factories.resource.local_db.helpers.request_snapshot import RequestSnapshot
+from stoobly_agent.app.proxy.record.join_request_service import join_request_from_request_response
+
+class SnapshotMigration():
+  _request: Request
+  _response: Response
+  _snapshot: RequestSnapshot
+
+  def __init__(self, snapshot: RequestSnapshot, log_event: LogEvent, log: Log = None):
+    self._event = log_event
+    self._log = log or Log()
+    self._request = snapshot.mitmproxy_request
+    self._response = snapshot.mitmproxy_response
+    self._snapshot = snapshot
+
+  @property
+  def request(self):
+    return self._request
+
+  @request.setter
+  def request(self, v: Request):
+    if not isinstance(v, Request):
+      raise TypeError('Invalid type.')
+    self._request = v
+
+  @property
+  def response(self):
+    return self._response
+
+  @response.setter
+  def response(self, v: Response):
+    if not isinstance(v, Response):
+      raise TypeError('Invalid type.')
+    self._response = v
+
+  @property
+  def snapshot(self):
+    return self._snapshot
+
+  @property
+  def uuid(self):
+    return self.snapshot.uuid
+
+  def delete(self, log: Log = None):
+    log = log or self._log
+    new_event = self._event.duplicate_as_delete()
+    log.append(str(new_event))
+
+  def save(self, log: Log = None):
+    log = log or self._log
+    request_uuid = self.snapshot.uuid
+    joined_request = join_request_from_request_response(self.request, self.response, id=request_uuid)
+    raw_request = joined_request.build()
+    self.snapshot.write_raw(raw_request)
+    new_event = self._event.duplicate()
+    log.append(str(new_event))
+    return new_event

stoobly_agent/app/models/adapters/orm/request/mitmproxy_adapter.py

@@ -15,4 +15,4 @@ class MitmproxyRequestAdapter():
     protocol = raw_http_request_adapter.protocol
     http_version = protocol.split('/')[1]
 
-    return PythonRequestMitmproxyRequestAdapter(http_version, python_request).adapt()
+    return PythonRequestMitmproxyRequestAdapter(python_request, http_version).adapt()

stoobly_agent/app/models/adapters/python/request/__init__.py

@@ -13,10 +13,10 @@ class PythonRequestAdapterFactory():
     self.__request = request
 
   def mitmproxy_request(self, http_version: str = 'HTTP/1.1'):
-    return MitmproxyRequestAdapter(http_version, self.__request).adapt()
+    return MitmproxyRequestAdapter(self.__request, http_version).adapt()
 
   def raw_request(self, http_version: str = 'HTTP/1.1'):
-    return RawRequestAdapter(http_version, self.__request).adapt()
+    return RawRequestAdapter(self.__request, http_version).adapt()
 
   def stoobly_request(self):
     return StooblyRequestAdapter(self.__request).adapt()

stoobly_agent/app/models/adapters/python/request/mitmproxy_adapter.py

@@ -8,7 +8,7 @@ from stoobly_agent.lib.utils.decode import decode
 
 class MitmproxyRequestAdapter():
 
-  def __init__(self, http_version: str, request: requests.Request):
+  def __init__(self, request: requests.Request, http_version: str = 'HTTP/1.1'):
     self.__http_version = http_version
     self.__request = request
 

stoobly_agent/app/models/adapters/python/request/raw_adapter.py

@@ -9,12 +9,12 @@ from .mitmproxy_adapter import MitmproxyRequestAdapter
 
 class RawRequestAdapter():
 
-  def __init__(self, http_version: str, request: requests.Request):
+  def __init__(self, request: requests.Request, http_version: str):
     self.__http_version = http_version
     self.__request = request
 
   def adapt(self):
-    mitmproxy_request = MitmproxyRequestAdapter(self.__http_version, self.__request).adapt()
+    mitmproxy_request = MitmproxyRequestAdapter(self.__request, self.__http_version).adapt()
     adapted_request = MitmproxyRequestFacade(mitmproxy_request)
     proxy_request = ProxyRequest(adapted_request)
     request_string = RequestString(proxy_request) 

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

@@ -19,19 +19,29 @@ class LocalDBHeaderAdapter(LocalDBAdapter):
     self.__request_orm = request_orm
 
   def create(self, request_id, **header: Header):
-    return self.update(request_id, **header)
+    name: str = header['name']
+    value = header['value']
 
-  def update(self, request_id, **header: Header):
-    request = self.__request_orm.find(request_id)
-    
-    if not request:
-      return self.__request_not_found()
+    headers = self.__headers(request_id)
+
+    for header in headers:
+      if header.lower() == name.lower(): 
+        return self.conflict(f"{name} already exists.")
 
+    headers[name] = value
+
+    LocalDBRequestAdapter(self.__request_orm).update(request_id, headers=headers)
+
+    return self.success({
+      'name': name,
+      'value': value,
+    })
+
+  def update(self, request_id, **header: Header):
     name = header['name']
     value = header['value']
 
-    request: requests.Request = RawHttpRequestAdapter(request.raw).to_request()
-    headers = request.headers
+    headers = self.__headers(request_id)
     headers[name] = value
 
     LocalDBRequestAdapter(self.__request_orm).update(request_id, headers=headers)
@@ -99,5 +109,15 @@ class LocalDBHeaderAdapter(LocalDBAdapter):
 
     return id
 
+  def __headers(self, request_id):
+    request = self.__request_orm.find(request_id)
+    
+    if not request:
+      return self.__request_not_found()
+
+    request: requests.Request = RawHttpRequestAdapter(request.raw).to_request()
+    headers = request.headers
+    return headers
+
   def __request_not_found(self):
     return self.not_found('Request not found')

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

@@ -82,6 +82,13 @@ class Log():
 
     return self.remove_processed_events(events, version_uuids)
 
+  @property
+  def resource_events(self) -> List[LogEvent]:
+    events = self.target_events
+    events = self.remove_dangling_events(events, events)
+
+    return list(filter(lambda e: e.action != DELETE_ACTION, events))
+
   @property
   def version(self):
     version_file_path = DataDir.instance().snapshosts_version_path

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

@@ -84,6 +84,16 @@ class LogEvent():
     toks[4] = str(int(time.time() * 1000))
     return LogEvent(COLUMN_DELIMITTER.join(toks))
 
+  def duplicate_as_delete(self):
+    event = self.duplicate()
+    event.action = DELETE_ACTION
+    return event
+
+  def duplicate_as_put(self):
+    event = self.duplicate()
+    event.action = PUT_ACTION
+    return event
+
   def is_request(self):
     return self.resource == REQUEST_RESOURCE
 

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

@@ -34,6 +34,26 @@ class RequestSnapshot(Snapshot):
     with open(request_file_path, 'rb') as fp:
       return fp.read()
 
+  @property
+  def mitmproxy_request(self):
+    from stoobly_agent.app.models.adapters.python.request.mitmproxy_adapter import MitmproxyRequestAdapter
+    return MitmproxyRequestAdapter(self.python_request).adapt()
+
+  @property
+  def mitmproxy_response(self):
+    from stoobly_agent.app.models.adapters.python.response.mitmproxy_adapter import MitmproxyResponseAdapter
+    return MitmproxyResponseAdapter(self.python_response).adapt()
+
+  @property
+  def python_request(self):
+    from stoobly_agent.app.models.adapters.raw_joined.request.python_adapter import PythonRequestAdapter
+    return PythonRequestAdapter(self.request).adapt()
+
+  @property
+  def python_response(self):
+    from stoobly_agent.app.models.adapters.raw_joined.response.python_adapter import PythonResponseAdapter
+    return PythonResponseAdapter(self.request).adapt()
+
   @property
   def path(self):
     dir_path = os.path.join(self.__requests_dir_path, self.uuid[0:2])

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

@@ -3,7 +3,7 @@ import os
 import pdb
 import shutil
 
-from typing import Callable
+from typing import Callable, List
 
 from stoobly_agent.lib.orm.scenario import Scenario
 
@@ -46,6 +46,13 @@ class ScenarioSnapshot(Snapshot):
   def metadata_path(self):
     return os.path.join(self.__scenarios_dir_path, self.uuid)
 
+  @property
+  def request_snapshots(self) -> List[RequestSnapshot]:
+    snapshots = []
+    handler = lambda request_snapshot: snapshots.append(request_snapshot)
+    self.iter_request_snapshots(handler)
+    return snapshots
+
   @property
   def requests(self):
     requests_file_path = self.requests_path

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

@@ -7,7 +7,7 @@ class Snapshot():
 
   def __init__(self, uuid: str):
     self.__uuid = uuid
-    self.__data_dir = DataDir.instance()
+    self.__data_dir: DataDir = DataDir.instance()
 
   @property
   def uuid(self):

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

@@ -36,6 +36,9 @@ class LocalDBAdapter():
   def not_found(self, d = 'Not Found'):
     return d, 404
 
+  def conflict(self, d = ''):
+    return d, 409
+
   def internal_error(self, d = 'Internal Error'):
     return d, 500
 

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

@@ -122,10 +122,10 @@ class LocalDBQueryParamAdapter(LocalDBAdapter):
     parsed_url = parsed_url._replace(query=urlencode(_query_params, True))
     request = LocalDBRequestAdapter(self.__request_orm).update(request_id, url=parsed_url.geturl())
 
-    return {
+    return self.success({
       'name': name,
       'value': value,
-    }
+    }) 
 
   def __decode_id(self, id: str):
     id = base64.b64decode(id)

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

@@ -193,13 +193,32 @@ class LocalDBRequestAdapter(LocalDBAdapter):
         }):
           return self.success(ORMToStooblyRequestTransformer(request, {}).transform())
     else:
-      if params.get('method'):
-        transformer.with_method(params['method'])
-
       if params.get('url'):
         transformer.with_url(params['url'])
         del params['url']
 
+      try:
+        if params.get('method'):
+          transformer.with_method(params['method'])
+
+        if params.get('scheme'):
+          transformer.with_scheme(params['scheme'])
+          # Do not delete scheme
+
+        if params.get('host'):
+          transformer.with_host(params['host'])
+          # Do not delete host
+
+        if params.get('port'):
+          transformer.with_port(params['port'])
+          # Do not delete port
+
+        if params.get('path'):
+          transformer.with_path(params['path'])
+          # Do not delete path
+      except ValueError as e:
+        return self.bad_request(str(e))
+
       if params.get('headers'):
         transformer.with_headers(params['headers'])
         del params['headers']

stoobly_agent/app/models/helpers/apply.py

@@ -218,15 +218,16 @@ class Apply():
 
     snapshot_requests = {}
 
-    raw_requests = snapshot.requests
-    for raw_request in raw_requests:
+    request_snapshots = snapshot.request_snapshots
+    for request_snapshot in request_snapshots:
+      raw_request = request_snapshot.request
+      
       toks = raw_request.split(REQUEST_STRING_CLRF.encode(), 1)
 
       if len(toks) != 2:
         return f"{snapshot.requests_path} contains an invalid request", 400
 
-      control = RequestStringControl(toks[0])
-      uuid = control.id
+      uuid = request_snapshot.uuid
       res, status = self.__put_request(uuid, raw_request, scenario_id=scenario['id'])
 
       if status != 200:

stoobly_agent/app/models/helpers/create_request_params_service.py

@@ -4,7 +4,7 @@ import requests
 from stoobly_agent.app.models.adapters import JoinedRequestAdapter, RawHttpRequestAdapter, RawHttpResponseAdapter
 from stoobly_agent.app.models.adapters.python import PythonRequestAdapterFactory, PythonResponseAdapterFactory
 
-from stoobly_agent.app.proxy.record.join_request_service import InterceptSettings, join_request, MitmproxyRequestFacade, MitmproxyResponseFacade
+from stoobly_agent.app.proxy.record.join_request_service import InterceptSettings, join_request_from_request_response 
 from stoobly_agent.app.settings import Settings
 from stoobly_agent.lib.logger import Logger
 
@@ -40,7 +40,9 @@ def build_params_from_python(request: requests.Request, response: requests.Respo
   mitmproxy_response = PythonResponseAdapterFactory(response).mitmproxy_response()
 
   intercept_settings = InterceptSettings(Settings.instance(), mitmproxy_request)
-  joined_request = join_request(MitmproxyRequestFacade(mitmproxy_request), MitmproxyResponseFacade(mitmproxy_response), intercept_settings)
+  joined_request = join_request_from_request_response(
+    mitmproxy_request, mitmproxy_response, intercept_settings=intercept_settings
+  )
 
   mitmproxy_flow_mock = MitmproxyFlowMock(mitmproxy_request, mitmproxy_response)
 

stoobly_agent/app/proxy/intercept_handler.py

@@ -29,11 +29,12 @@ def request(flow: MitmproxyHTTPFlow):
     __patch_cookie(request)
 
     intercept_settings = InterceptSettings(Settings.instance(), request)
-
     if not intercept_settings.active:
         return
 
     __disable_web_cache(request)
+    __disable_content_encoding(request)
+
     __intercept_hook(lifecycle_hooks.BEFORE_REQUEST, flow, intercept_settings)
 
     active_mode = intercept_settings.mode
@@ -86,10 +87,13 @@ def response(flow: MitmproxyHTTPFlow):
 ### PRIVATE
 
 # Prevent 304 status
-# Because this header will get recorded, should add during mocking as well in the case where headers are used for matching
 def __disable_web_cache(request: MitmproxyRequest) -> None:
     request.headers['Cache-Control'] = 'no-cache, no-store'
 
+# Disable response body returning as encoded e.g. gzip
+def __disable_content_encoding(request: MitmproxyRequest) -> None:
+    request.headers['Accept-Encoding'] = 'identity'
+
 # Fix issue where multi-value cookies become comma separated
 def __patch_cookie(request: MitmproxyRequest):
     header_name = 'cookie'

stoobly_agent/app/proxy/record/join_request_service.py

@@ -1,31 +1,43 @@
 import pdb
 
-from mitmproxy.http import HTTPFlow as MitmproxyHTTPFlow
-from typing import List
+from mitmproxy.http import HTTPFlow as MitmproxyHTTPFlow, Request as MitmproxyRequest, Response as MitmproxyResponse
+from typing import TypedDict
 
 from stoobly_agent.app.proxy.intercept_settings import InterceptSettings
-from stoobly_agent.app.settings.rewrite_rule import RewriteRule
 
 from ..mitmproxy.request_facade import MitmproxyRequestFacade
 from ..mitmproxy.response_facade import MitmproxyResponseFacade
-from ..utils.rewrite import rewrite_request_response
 from .joined_request import JoinedRequest
 from .proxy_request import ProxyRequest
 
+class JoinRequestOptions(TypedDict):
+    id: str
+    intercept_settings: InterceptSettings
+
 def join_request(
-    adapted_request: MitmproxyRequestFacade, adapted_response: MitmproxyResponseFacade, intercept_settings: InterceptSettings
+    adapted_request: MitmproxyRequestFacade, adapted_response: MitmproxyResponseFacade, **options: JoinRequestOptions
 ) -> JoinedRequest:
+    intercept_settings: InterceptSettings = options.get('intercept_settings')
+
     # Decorate request with service_url
-    upstream_url = intercept_settings.upstream_url
+    upstream_url = intercept_settings.upstream_url if intercept_settings else None
     proxy_request = ProxyRequest(adapted_request, upstream_url)
 
+    if options.get('id'):
+        proxy_request.id = options['id']
+
     # Create JoinedRequest
     return JoinedRequest(proxy_request).with_response(adapted_response)
 
 def join_request_from_flow(
-    flow: MitmproxyHTTPFlow, intercept_settings: InterceptSettings
+    flow: MitmproxyHTTPFlow, **options: JoinRequestOptions
 ) -> JoinedRequest:
-    request = MitmproxyRequestFacade(flow.request)
-    response = MitmproxyResponseFacade(flow.response)
+    return join_request_from_request_response(flow.request, flow.response, **options)
+
+def join_request_from_request_response(
+    request: MitmproxyRequest, response: MitmproxyResponse, **options: JoinRequestOptions
+):
+    request = MitmproxyRequestFacade(request)
+    response = MitmproxyResponseFacade(response)
 
-    return join_request(request, response, intercept_settings)
+    return join_request(request, response, **options)

stoobly_agent/app/proxy/record/proxy_request.py

@@ -2,10 +2,19 @@ from ..mitmproxy.request_facade import MitmproxyRequestFacade
 
 class ProxyRequest:
     def __init__(self, request: MitmproxyRequestFacade, upstream_url: str = None):
+        self._id = None
         self.request = request
 
         self.upstream_url = upstream_url
 
+    @property
+    def id(self):
+        return self._id
+
+    @id.setter
+    def id(self, v):
+        self._id = v
+
     def url(self):
         url = self.request.url
 

stoobly_agent/app/proxy/record/request_string.py

@@ -15,6 +15,7 @@ class RequestString:
     ENCODING = 'utf-8'
 
     __current_time = None
+    request_id = None
 
     def __init__(self, proxy_request: ProxyRequest):
         self.__current_time = self.__get_current_time()
@@ -29,7 +30,9 @@ class RequestString:
             self.__headers()
             self.__body()
 
-        self.request_id = self.__generate_request_id()
+            self.request_id = proxy_request.id
+
+        self.request_id = self.request_id or self.__generate_request_id()
 
     def get(self, **kwargs):
         if kwargs.get('control'):

stoobly_agent/app/proxy/record/upload_request_service.py

@@ -4,6 +4,7 @@ import pdb
 import time
 import tempfile
 
+from copy import deepcopy
 from mitmproxy.http import HTTPFlow as MitmproxyHTTPFlow
 from mitmproxy.http import Request as MitmproxyRequest
 
@@ -50,22 +51,23 @@ def upload_request(
 ):
     Logger.instance(LOG_ID).info(f"{bcolors.OKCYAN}Recording{bcolors.ENDC} {flow.request.url}")
 
-    joined_request = join_request_from_flow(flow, intercept_settings)
+    flow_copy = deepcopy(flow) # When applying modifications we don't want to persist them in the response
+    joined_request = join_request_from_flow(flow_copy, intercept_settings=intercept_settings)
 
     project_key = intercept_settings.project_key
     scenario_key = intercept_settings.scenario_key
     body_params = __build_body_params(
         project_key,
-        joined_request, 
+        joined_request,
         flow=flow,
         scenario_key=scenario_key
     )
 
     # If request_origin is WEB, then we are in proxy
     # This means that we have access to Cache singleton and do not need send a request to update the status
-    sync = intercept_settings.request_origin == request_origin.WEB 
+    sync = intercept_settings.request_origin == request_origin.WEB
     res = __upload_request_with_body_params(request_model, body_params, sync)
- 
+
     if intercept_settings.settings.is_debug():
         file_path = __debug_request(flow.request, joined_request.build())
         Logger.instance(LOG_ID).debug(f"Writing request to {file_path}")
@@ -142,4 +144,4 @@ def __debug_request(request: MitmproxyRequest, raw_requests: bytes):
     with open(file_path, 'wb') as f:
         f.write(raw_requests)
 
-    return file_path
+    return file_path

stoobly_agent/app/proxy/test/helpers/upload_test_service.py

@@ -43,7 +43,7 @@ def upload_test(
   flow: MitmproxyHTTPFlow, 
   **kwargs: UploadTestData
 ) -> Response:
-    joined_request = join_request_from_flow(flow, intercept_settings)
+    joined_request = join_request_from_flow(flow, intercept_settings=intercept_settings)
 
     Logger.instance(LOG_ID).info(f"{bcolors.OKCYAN}Uploading{bcolors.ENDC} test results for {joined_request.proxy_request.url()}")
 

stoobly_agent/config/constants/lifecycle_hooks.py

@@ -3,6 +3,7 @@ AFTER_RECORD = 'handle_after_record'
 AFTER_REPLAY = 'handle_after_replay'
 AFTER_TEST = 'handle_after_test'
 BEFORE_MOCK = 'handle_before_mock'
+BEFORE_MIGRATE = 'handle_before_migrate'
 BEFORE_RECORD = 'handle_before_record'
 BEFORE_REPLAY = 'handle_before_replay'
 BEFORE_REQUEST = 'handle_before_request'