cavisson-pythonagent 0.0.1__py3-none-any.whl

pythonagent/agent/probes/mongodb/pymongo.py

@@ -0,0 +1,286 @@
+from __future__ import unicode_literals
+#import sys
+import inspect
+
+from ..base import ExitCallInterceptor
+#import abc
+#import agent
+#from agent.internal.proxy import *
+from pythonagent.lang import str
+from pythonagent.utils import generate_flow_path_id, get_utf8_bytes, get_current_timestamp_in_us, get_validated_duration
+
+
+EXIT_CUSTOM = 1
+EXIT_SUBTYPE_MONGODB = 'Mongo DB'
+funcName = 'pymongo.monitoring.CommandListener.started'
+
+COLLECTION_METHODS_FOR_HAVOC = [
+    'find_one', 'find_one_and_update', 'find_one_and_replace', 'find_one_and_delete',
+    'insert_one', 'insert_many',
+    'update_one', 'update_many', 'replace_one',
+    'delete_one', 'delete_many',
+    'aggregate', 'count_documents', 'bulk_write',
+]
+
+
+def _get_collection_host_info(col):
+    """Extract (host, port, db_name) from a Collection object. Never raises.
+
+    Sync-only helper for pymongo.MongoClient/Collection, where `.address` and
+    `.nodes` are plain properties. Do NOT use this for AsyncCollection — use
+    _get_collection_host_info_async instead, since AsyncMongoClient.address
+    is a coroutine and awaiting it here isn't possible from a sync function.
+    """
+    try:
+        client = col.database.client
+        address = getattr(client, 'address', None)
+        if not address:
+            nodes = getattr(client, 'nodes', None)
+            address = next(iter(nodes), None) if nodes else None
+        if address:
+            host, port = address
+        else:
+            host = getattr(client, 'HOST', 'localhost')
+            port = getattr(client, 'PORT', 27017)
+        db_name = col.database.name
+        return host, port, db_name
+    except Exception:
+        return 'unknown', 27017, 'unknown'
+
+
+async def _get_collection_host_info_async(col):
+    """Async-safe (host, port, db_name) extraction for pymongo.asynchronous.AsyncCollection.
+
+    On the native async driver, AsyncMongoClient.address (and .nodes) are
+    coroutines and must be awaited. Calling the sync _get_collection_host_info
+    against an AsyncCollection creates-but-never-awaits that coroutine, which
+    silently falls into the except-fallback AND leaks an unawaited coroutine
+    object — that's the source of:
+
+        RuntimeWarning: coroutine 'AsyncMongoClient.address' was never awaited
+
+    This variant detects the awaitable and awaits it properly instead.
+    """
+    try:
+        client = col.database.client
+        address = getattr(client, 'address', None)
+        if inspect.isawaitable(address):
+            address = await address
+        if not address:
+            nodes = getattr(client, 'nodes', None)
+            if inspect.isawaitable(nodes):
+                nodes = await nodes
+            address = next(iter(nodes), None) if nodes else None
+        if address:
+            host, port = address
+        else:
+            host = getattr(client, 'HOST', 'localhost')
+            port = getattr(client, 'PORT', 27017)
+        db_name = col.database.name
+        return host, port, db_name
+    except Exception:
+        return 'unknown', 27017, 'unknown'
+
+
+def _patch_collection_for_havoc(agent, collection_cls):
+    """Monkey-patch Collection CRUD methods so NDHavocException propagates to the caller.
+
+    The CommandListener monitoring API cannot do this — PyMongo swallows all exceptions
+    raised inside monitoring callbacks, so NDHavocException never reaches the application.
+    Wrapping Collection methods directly ensures the exception is in the call stack.
+    """
+
+    def make_wrapper(original_method, mname):
+        def havoc_wrapper(col_self, *args, **kwargs):
+            bt = agent.get_current_bt()
+            if bt:
+                host, port, db_name = _get_collection_host_info(col_self)
+                host_str = "NA|{}|{}|NA|MONGODB|{}|NA|NA|NA|NA|NA".format(host, port, db_name)
+                try:
+                    from pythonagent.agent.probes.havoc.havoc_manager import NDHavocException
+                    backend_name_value = agent.setting_backend_name_otel(host_str, "DB")
+                    if backend_name_value:
+                        backend_name = (
+                            backend_name_value.decode()
+                            if isinstance(backend_name_value, bytes)
+                            else backend_name_value
+                        )
+                    else:
+                        backend_name = None
+                    ctx = agent.get_transaction_context()
+                    btname = ctx.btname if ctx else None
+                    req_url = ctx.url_path if ctx else None
+                    agent.havoc_monitor.apply_outbound_service_failure(host_str, backend_name,req_url,btname)
+                except Exception as e:
+                    from pythonagent.agent.probes.havoc.havoc_manager import NDHavocException
+                    if isinstance(e, NDHavocException):
+                        agent.logger.info("MongoDB Collection havoc triggered for {}: {}".format(mname, e))
+                        agent.set_current_status_code(503)
+                        db_start_time = get_current_timestamp_in_us()
+                        ip_handle = agent.db_call_begin(funcName, bt, str(host), mname, None, db_start_time)
+                        agent.db_call_end(bt, ip_handle, 0, db_start_time, funcName, host_str, 503)
+                        raise e
+                    agent.logger.debug("Non-havoc exception in MongoDB havoc check: {}".format(e))
+            return original_method(col_self, *args, **kwargs)
+
+        havoc_wrapper._pythonagent_intercepted = True
+        return havoc_wrapper
+
+    for method_name in COLLECTION_METHODS_FOR_HAVOC:
+        original = getattr(collection_cls, method_name, None)
+        if original is None:
+            continue
+        if getattr(original, '_pythonagent_intercepted', False):
+            continue
+        setattr(collection_cls, method_name, make_wrapper(original, method_name))
+
+
+def _patch_async_collection_for_havoc(agent, async_collection_cls):
+    """Async mirror of _patch_collection_for_havoc for pymongo.asynchronous.AsyncCollection.
+
+    AsyncCollection methods are coroutines (async def). The wrapper must also be
+    async def and must await the original method so the call chain is preserved.
+    Havoc checks are synchronous and safe to call from async context; time.sleep
+    in apply_outbound_service_delay intentionally blocks the event loop for the delay.
+    """
+
+    def make_async_wrapper(original_method, mname):
+        async def havoc_wrapper(col_self, *args, **kwargs):
+            bt = agent.get_current_bt()
+            if bt:
+                host, port, db_name = await _get_collection_host_info_async(col_self)
+                host_str = "NA|{}|{}|NA|MONGODB|{}|NA|NA|NA|NA|NA".format(host, port, db_name)
+                try:
+                    from pythonagent.agent.probes.havoc.havoc_manager import NDHavocException
+                    backend_name_value = agent.setting_backend_name_otel(host_str, "DB")
+                    if backend_name_value:
+                        backend_name = (
+                            backend_name_value.decode()
+                            if isinstance(backend_name_value, bytes)
+                            else backend_name_value
+                        )
+                    else:
+                        backend_name = None
+                    ctx = agent.get_transaction_context()
+                    btname = ctx.btname if ctx else None
+                    req_url = ctx.url_path if ctx else None
+                    agent.havoc_monitor.apply_outbound_service_failure(host_str, backend_name,req_url,btname)
+                except Exception as e:
+                    from pythonagent.agent.probes.havoc.havoc_manager import NDHavocException
+                    if isinstance(e, NDHavocException):
+                        agent.logger.info("MongoDB AsyncCollection havoc triggered for {}: {}".format(mname, e))
+                        agent.set_current_status_code(503)
+                        db_start_time = get_current_timestamp_in_us()
+                        ip_handle = agent.db_call_begin(funcName, bt, str(host), mname, None, db_start_time)
+                        agent.db_call_end(bt, ip_handle, 0, db_start_time, funcName, host_str, 503)
+                        raise e
+                    agent.logger.debug("Non-havoc exception in MongoDB async havoc check: {}".format(e))
+            return await original_method(col_self, *args, **kwargs)
+
+        havoc_wrapper._pythonagent_intercepted = True
+        return havoc_wrapper
+
+    for method_name in COLLECTION_METHODS_FOR_HAVOC:
+        original = getattr(async_collection_cls, method_name, None)
+        if original is None:
+            continue
+        if getattr(original, '_pythonagent_intercepted', False):
+            continue
+        setattr(async_collection_cls, method_name, make_async_wrapper(original, method_name))
+
+
+def intercept_pymongo(agent, mod):
+    class ExitCallListener(mod.monitoring.CommandListener):
+        backend_name_format_string = '{HOST}:{PORT} - {DATABASE}'
+
+        def __init__(self):
+            self.interceptor = ExitCallInterceptor(agent, None)
+            self.entry_point_fqm = None
+            self.exit_call_map = {}
+
+        def get_backend(self, connection_id, database_name):
+            host, port = connection_id
+            backend_properties = {
+                'HOST': host,
+                'PORT': str(port),
+                'DATABASE': database_name,
+            }
+            agent.logger.info("Modulenameintercept_pymongo  class || backend_properties is :{0}".format(backend_properties))
+            return backend_properties
+
+        def started(self, event):
+            try:
+                db_start_time = 0
+                host_, port_ = event.connection_id
+                # format :- "PROTOCOL|HOST|PORT|URL|DBPRODNAME|DBNAME|DBPRODVER|DRIVERNAME|DRIVERVER|USERNAME|SERVICENAME"
+                host_name_str = f"NA|{host_}|{port_}|NA|MONGODB|{event.database_name}|NA|NA|NA|NA|NA"
+                if not self.entry_point_fqm:
+                    mName = mod.monitoring.CommandListener.__module__ + "." + mod.monitoring.CommandListener.__name__ + ".started"
+                    self.entry_point_fqm = mName
+
+                bt = self.interceptor.bt
+                if not bt:
+                    bt = agent.get_current_bt()
+                if bt:
+                    backend = self.get_backend(event.connection_id, event.database_name)
+                    if backend:
+                        if backend:
+                            db_start_time = get_current_timestamp_in_us()
+                            agent.logger.debug("db_call_begin for mongodb bt {} host {} query {} query_params {} db_start_time {}".format(bt, backend["HOST"], str(event.command_name), None, db_start_time))
+                            exit_call = self.interceptor.db_call_begin(bt, backend["HOST"], str(event.command_name), None, db_start_time, self.entry_point_fqm, host_name_str)
+                        self.exit_call_map[event.operation_id] = (exit_call, host_name_str, db_start_time)
+            except Exception as e:
+                from pythonagent.agent.probes.havoc.havoc_manager import NDHavocException
+                if isinstance(e, NDHavocException):
+                    raise e
+                agent.logger.exception("Exception occurred in pymongo started {}".format(e))
+
+        def succeeded(self, event):
+            try:
+                if not self.interceptor.bt:
+                    return
+
+                db_end_time = get_current_timestamp_in_us()
+                ip_handle, host_str, db_start_time = self.exit_call_map.pop(event.operation_id, None)
+                duration = get_validated_duration(db_start_time, db_end_time, "MongoDB")
+                status_code = 200
+
+                agent.logger.debug("db_call_end for mongodb succeeded bt {} ip_handle {} duration {} db_start_time {}".format(self.interceptor.bt, ip_handle, duration, db_start_time))
+                self.interceptor.db_call_end(self.interceptor.bt, ip_handle, duration, db_start_time, self.entry_point_fqm, host_str, status_code)
+            except Exception as e:
+                self.agent.logger.exception("Exception occurred in pymongo succeeded {}".format(e))
+
+        def failed(self, event):
+            try:
+                if not self.interceptor.bt:
+                    return
+
+                db_end_time = get_current_timestamp_in_us()
+                ip_handle, host_str, db_start_time = self.exit_call_map.pop(event.operation_id, None)
+                duration = get_validated_duration(db_start_time, db_end_time, "MongoDB")
+                status_code = 500
+                agent.logger.debug("db_call_end for mongodb failed bt {} ip_handle {} duration {} db_start_time {}".format(self.interceptor.bt, ip_handle, duration, db_start_time))
+                self.interceptor.db_call_end(self.interceptor.bt, ip_handle, duration, db_start_time, self.entry_point_fqm, host_str, status_code)
+            except Exception as e:
+                self.agent.logger.exception("Exception occurred in pymongo failed {}".format(e))
+
+    mod.monitoring.register(ExitCallListener())
+
+    # Patch Collection CRUD methods so NDHavocException propagates to the caller.
+    # CommandListener cannot do this — PyMongo swallows exceptions from monitoring callbacks.
+    try:
+        import importlib
+        _patch_collection_for_havoc(agent, importlib.import_module('pymongo.collection').Collection)
+    except Exception as e:
+        agent.logger.warning("Could not patch pymongo Collection for havoc injection: {}".format(e))
+
+    # Patch AsyncCollection for PyMongo 4.7+ native async API.
+    # Gracefully skipped on older versions where pymongo.asynchronous does not exist.
+    try:
+        import importlib
+        _patch_async_collection_for_havoc(
+            agent,
+            importlib.import_module('pymongo.asynchronous.collection').AsyncCollection
+        )
+    except Exception as e:
+        agent.logger.warning("Could not patch pymongo AsyncCollection for havoc injection: {}".format(e))

pythonagent/agent/probes/openai/__init__.py

@@ -0,0 +1,3 @@
+from .openai import intercept_openai
+
+__all__ = ['intercept_openai']