cavisson-pythonagent 0.0.1__py3-none-any.whl

pythonagent/agent/probes/http/__init__.py

@@ -0,0 +1,49 @@
+"""Base class for HTTP connection interceptors.
+
+"""
+
+from __future__ import unicode_literals
+
+from ..base import ExitCallInterceptor
+from pythonagent.lang import str, urlparse
+
+
+EXIT_HTTP =0
+EXIT_SUBTYPE_HTTP = 'HTTP'
+
+class HTTPConnectionInterceptor(ExitCallInterceptor):
+    # If the library you are intercepting has an HTTPSConnection class which
+    # does not subclass httplib.HTTPSConnection, add it to this set.
+    https_connection_classes = set()
+    backend_name_format_string = '%s://{HOST}:{PORT}{URL}?{QUERY STRING}'
+
+    @classmethod
+    def _request_is_https(cls, connection):
+        #print("INSIDE IS REQUEST HTTPS")
+        if connection.port == 443:
+            return True
+        return isinstance(connection, tuple(cls.https_connection_classes))
+
+    def get_backend(self, host, port, scheme, url):
+        #self.agent.logger.info('Modulename: HTTPConnectionInterceptor class get_backend function')
+        parsed_url = urlparse(url)
+        backend_properties = {
+            'HOST': host,
+            'PORT': str(port),
+            'URL': parsed_url.path,
+            'QUERY STRING': parsed_url.query,
+        }
+        return self.agent.backend_registry.get_backend(EXIT_HTTP, EXIT_SUBTYPE_HTTP, backend_properties,
+                                                       self.backend_name_format_string % scheme)
+
+
+from .httplib import intercept_httplib
+from .urllib3 import intercept_urllib3
+from .requests import intercept_requests
+from .tornado_httpclient import intercept_tornado_httpclient
+from .boto import intercept_boto
+from .aiohttp_client import intercept_aiohttp_client
+from .httpx_client import intercept_httpx_client
+
+__all__ = ['intercept_httplib', 'intercept_urllib3', 'intercept_requests', 'intercept_tornado_httpclient', 'intercept_boto', 'intercept_aiohttp_client', 'intercept_httpx_client']
+

pythonagent/agent/probes/http/aiohttp_client.py

@@ -0,0 +1,59 @@
+from ..base import ExitCallInterceptor
+
+from pythonagent.utils import get_current_timestamp_in_us, get_validated_duration, generate_callout_id
+from urllib.parse import urlparse
+
+
+class AsyncioHTTPClientInterceptor(ExitCallInterceptor):
+
+    async def wrap_request(self, _request, *args, **kwargs):
+        bt = None
+        exit_call = None
+        start_time = 0
+        try:
+            unique_id = generate_callout_id()
+            full_url = args[2]
+            if full_url:
+                parsed_url = urlparse(full_url)
+                host = parsed_url.hostname
+                port = parsed_url.port or ('443' if parsed_url.scheme == 'https' else '80')
+                url = parsed_url.path
+                host_port = "NA|" + str(host) + "|" + str(port) + "|" + url
+            else:
+                host_port = "NA|default_host|default_port|default_url"
+
+            mName = _request.__module__ + "." + _request.__qualname__
+            headers_value = self.agent.get_nd_header(self.bt, host_port, unique_id).decode()
+
+            client_session_object = args[0]
+
+            if client_session_object._default_headers:
+                client_session_object._default_headers["CavNDFPInstance"] = headers_value
+            else:
+                client_session_object._default_headers = {"CavNDFPInstance": headers_value}
+
+            bt = self.bt
+            start_time = get_current_timestamp_in_us()
+            exit_call = self.http_call_begin(bt, host_port, url, mName, start_time, host_port, unique_id)
+        except Exception as e:
+            from pythonagent.agent.probes.havoc.havoc_manager import NDHavocException
+            if isinstance(e, NDHavocException):
+                raise e
+            self.agent.logger.exception("Exception in http call begin in aioHttp Client  {}".format(e))
+
+        result = await _request(*args, **kwargs)
+
+        try:
+            end_time = get_current_timestamp_in_us()
+            duration = get_validated_duration(start_time, end_time, "AIOHTTP Client")
+            self.http_call_end(self.bt, exit_call, mName, result.status, duration, start_time, host_port, unique_id)
+
+        except Exception as e:
+            self.agent.logger.exception("Exception in http call end in aioHttp Client  {}".format(e))
+        return result
+
+
+
+def intercept_aiohttp_client(agent, mod):
+    agent.logger.warning("Instrument module: aiohttp.client{}".format(", mod: {}".format(mod) if mod else ""))
+    return AsyncioHTTPClientInterceptor(agent, mod.ClientSession).attach('_request', patched_method_name='wrap_request')

pythonagent/agent/probes/http/boto.py

@@ -0,0 +1,12 @@
+
+"""Intercept boto to ensure that HTTPS is reported correctly.
+
+"""
+
+from __future__ import unicode_literals
+
+from . import HTTPConnectionInterceptor
+
+
+def intercept_boto(agent, mod):
+    HTTPConnectionInterceptor.https_connection_classes.add(mod.CertValidatingHTTPSConnection)

pythonagent/agent/probes/http/httplib.py

@@ -0,0 +1,110 @@
+
+
+"""Interceptor for httplib/http.client.
+
+"""
+
+from __future__ import unicode_literals
+from . import HTTPConnectionInterceptor
+from pythonagent.utils import get_current_timestamp_in_us, get_validated_duration, generate_callout_id
+
+
+class HttplibConnectionInterceptor(HTTPConnectionInterceptor):
+    def _putrequest(self, putrequest, connection, method, url, *args, **kwargs):
+        exit_call = None
+        with self.log_exceptions():
+            ctx = self.agent.get_transaction_context()
+            if ctx:
+                ctx.http_request_type = method  # GET, PUT, POST or DELETE
+            bt = self.bt
+            host_port = "NA|" + str(connection.host) + "|" + str(connection.port) + "|" + url
+            mName = str(putrequest.__module__ + "." + putrequest.__qualname__)
+            if ctx:
+                ctx.entry_point_fqm = mName
+                ctx.backend_header = host_port
+                ctx.callout_id = generate_callout_id()
+            current_time = get_current_timestamp_in_us()
+            if ctx:
+                ctx.tier_callout_start_time = current_time
+            exit_call = self.http_call_begin(bt, host_port, url, mName, current_time, host_port, ctx.callout_id if ctx else 0)
+            if bt:
+                scheme = 'https' if self._request_is_https(connection) else 'http'
+                connection._pythonagent_exit_call = exit_call
+        return putrequest(connection, method, url, pythonagent_exit_call=exit_call, *args, **kwargs)
+
+    def _endheaders(self, endheaders, connection, *args, **kwargs):
+        exit_call = None
+        with self.log_exceptions():
+            bt = self.bt
+            if self.agent.cav_env == "NATIVE":
+                prop_header = self.agent.sdk_get_propagation_header(bt)
+            else:
+                prop_header = None
+            if prop_header:
+                prop_header = prop_header.decode()
+                prop_header = str(prop_header)
+                prop_header = prop_header.split('\n')
+                for pr_header in prop_header:
+                    key,value = pr_header.split(':')
+                    connection.putheader(key,value)
+
+            if self.agent.cav_env == "NATIVE":
+                ctx = self.agent.get_transaction_context()
+                if ctx:
+                    backend_name = ctx.backend_header
+                    callout_id = ctx.callout_id
+                    headers_value = self.agent.get_nd_header(self.bt, backend_name, callout_id)
+                    if headers_value:
+                        connection.putheader('CavNDFPInstance', headers_value.decode())
+
+            exit_call = getattr(connection, '_pythonagent_exit_call', None)
+            #self.agent.logger.info("Modulename: HttplibConnectionInterceptor class inside _endheaders exit_call is :{0}".format(exit_call))
+
+            header = self.make_correlation_header(exit_call)
+            if header is not None:
+                connection.putheader(*header)
+                #self.agent.logger.debug('Added correlation header to HTTP request: %s, %s' % header)
+        return endheaders(connection, pythonagent_exit_call=exit_call, *args, **kwargs)
+
+    def _getresponse(self, getresponse, connection, *args, **kwargs):
+        # CORE-40945 Catch TypeError as a special case for Python 2.6 and call getresponse with just the HTTPConnection instance.
+        exit_call = None
+        try:
+            exit_call = getattr(connection, '_pythonagent_exit_call', None)
+            bt = self.bt
+        except Exception as e:
+           self.agent.logger.exception("Error in _getresponse begin part: %s", e)
+
+        try:
+            with self.end_exit_call_and_reraise_on_exception(exit_call, ignored_exceptions=(TypeError,)):
+                response = getresponse(connection, *args, **kwargs)
+        except TypeError:
+            with self.end_exit_call_and_reraise_on_exception(exit_call):
+                response = getresponse(connection)
+
+        try:
+            current_time = get_current_timestamp_in_us()
+            ctx = self.agent.get_transaction_context()
+            start_time = ctx.tier_callout_start_time if ctx else 0
+            duration = get_validated_duration(start_time, current_time, "HTTP")
+            if ctx:
+                ctx.tier_callout_start_time = 0
+            method_name = None
+            self.http_call_end(self.bt, exit_call, method_name, response.status, duration, start_time, ctx.backend_header if ctx else None, ctx.callout_id if ctx else 0)
+        except Exception as e:
+            self.agent.logger.exception("Error in _getresponse end part: %s", e)
+
+        try:
+            del connection._pythonagent_exit_call
+        except AttributeError:
+            pass
+        return response
+
+
+def intercept_httplib(agent, mod):
+    #print("intercept_httplib")
+    agent.logger.warning("Instrument module: httplib/http.client{}".format(", mod: {}".format(mod) if mod else ""))
+    HTTPConnectionInterceptor.https_connection_classes.add(mod.HTTPSConnection)
+    interceptor = HttplibConnectionInterceptor(agent, mod.HTTPConnection)
+    interceptor.attach(['putrequest', 'endheaders'])
+    interceptor.attach('getresponse', wrapper_func=None)   # CORE-40945 Do not wrap getresponse in the default wrapper.

pythonagent/agent/probes/http/httpx_client.py

@@ -0,0 +1,116 @@
+from ..base import ExitCallInterceptor
+
+from pythonagent.utils import get_current_timestamp_in_us, get_validated_duration, generate_callout_id
+from urllib.parse import urlparse
+
+
+class HTTPXAsyncClientInterceptor(ExitCallInterceptor):
+
+    async def wrap_send(self, send, *args, **kwargs):
+        bt = None
+        exit_call = None
+        start_time = 0
+
+        try:
+            unique_id = generate_callout_id()
+            url = args[1].url
+            full_url = str(url)
+            if full_url:
+                parsed_url = urlparse(full_url)
+                host = parsed_url.hostname
+                port = parsed_url.port or ('443' if parsed_url.scheme == 'https' else '80')
+                url = parsed_url.path or "/"
+                if parsed_url.query:
+                    url = url + "?" + parsed_url.query
+                host_port = "NA|" + host + "|" + str(port) + "|" + url
+            else:
+                host_port = "NA|default_host|default_port|default_url"
+
+            bt = self.bt
+            mName = send.__module__ + "." + send.__qualname__
+            headers_value = self.agent.get_nd_header(self.bt, host_port, unique_id).decode()
+
+            client_session_object = args[0]
+
+            if client_session_object._headers:
+                client_session_object._headers["CavNDFPInstance"] = headers_value
+            else:
+                client_session_object._headers = {"CavNDFPInstance": headers_value}
+
+
+            start_time = get_current_timestamp_in_us()
+            exit_call = self.http_call_begin(bt, host_port, url, mName, start_time, host_port, unique_id)
+        except Exception as e:
+            from pythonagent.agent.probes.havoc.havoc_manager import NDHavocException
+            if isinstance(e, NDHavocException):
+                raise e
+            self.agent.logger.exception("Exception in http call begin in httpx_async Client  {}".format(e))
+
+        result = await send(*args, **kwargs)
+
+        try:
+            current_time = get_current_timestamp_in_us()
+            duration = get_validated_duration(start_time, current_time, "httpx_async Client")
+            self.http_call_end(self.bt, exit_call, mName, result.status_code, duration, start_time, host_port, unique_id)
+
+        except Exception as e:
+            self.agent.logger.exception("Exception in http call end in httpx_async Client  {}".format(e))
+
+        return result
+
+
+class HTTPXClientInterceptor(ExitCallInterceptor):
+    def wrap_send(self, send, *args, **kwargs):
+        bt = None
+        exit_call = None
+        start_time = 0
+
+        try:
+            unique_id = generate_callout_id()
+            url = args[1].url
+            full_url = str(url)
+            if full_url:
+                parsed_url = urlparse(full_url)
+                host = parsed_url.hostname
+                port = parsed_url.port or ('443' if parsed_url.scheme == 'https' else '80')
+                url = parsed_url.path or "/"
+                if parsed_url.query:
+                    url = url + "?" + parsed_url.query
+                host_port = "NA|" + host + "|" + str(port) + "|" + url
+            else:
+                host_port = "NA|default_host|default_port|default_url"
+            bt = self.bt
+            mName = send.__module__ + "." + send.__qualname__
+
+            headers_value = self.agent.get_nd_header(self.bt, host_port, unique_id).decode()
+            client_session_object = args[0]
+            if client_session_object._headers:
+                client_session_object._headers["CavNDFPInstance"] = headers_value
+            else:
+                client_session_object._headers = {"CavNDFPInstance": headers_value}
+
+            start_time = get_current_timestamp_in_us()
+            exit_call = self.http_call_begin(bt, host_port, url, mName, start_time, host_port, unique_id)
+        except Exception as e:
+            from pythonagent.agent.probes.havoc.havoc_manager import NDHavocException
+            if isinstance(e, NDHavocException):
+                raise e
+            self.agent.logger.exception("Exception in http call begin in httpx_sync Client  {}".format(e))
+
+
+        result = send(*args, **kwargs)
+
+        try:
+            current_time = get_current_timestamp_in_us()
+            duration = get_validated_duration(start_time, current_time, "httpx_sync Client")
+            self.http_call_end(self.bt, exit_call, mName, result.status_code, duration, start_time, host_port, unique_id)
+        except Exception as e:
+            self.agent.logger.exception("Exception in http call begin in httpx_sync Client  {}".format(e))
+
+        return result
+
+
+def intercept_httpx_client(agent, mod):
+    agent.logger.warning("Instrument module: httpx{}".format(", mod: {}".format(mod) if mod else ""))
+    HTTPXClientInterceptor(agent, mod.Client).attach('send', patched_method_name='wrap_send')
+    HTTPXAsyncClientInterceptor(agent, mod.AsyncClient).attach('send', patched_method_name='wrap_send')

pythonagent/agent/probes/http/requests.py

@@ -0,0 +1,15 @@
+
+"""Intercept requests to ensure that HTTPS is reported correctly.
+
+"""
+
+from __future__ import unicode_literals
+
+from .urllib3 import intercept_urllib3
+
+
+def intercept_requests(agent, mod):
+    #print("intercept_requests")
+    agent.logger.warning("Instrument module: requests{}".format(", mod: {}".format(mod) if mod else ""))
+    # requests ships with its own version of urllib3, so we need to manually intercept it.
+    intercept_urllib3(agent, mod.packages.urllib3)

pythonagent/agent/probes/http/tornado_httpclient.py

@@ -0,0 +1,85 @@
+from __future__ import unicode_literals
+import functools
+from enum import unique
+from logging import exception
+
+from pythonagent.lang import get_args, urlparse
+
+from . import HTTPConnectionInterceptor
+from pythonagent.utils import get_current_timestamp_in_us, get_validated_duration, generate_callout_id
+
+
+class TornadoAsyncHTTPClientInterceptor(HTTPConnectionInterceptor):
+
+    def _fetch(self, fetch, *args, **kwargs):
+        bt = self.bt
+        if not bt:
+            return fetch(*args, **kwargs)
+
+        exit_call = None
+        start_time = 0
+
+        try:
+            unique_id = generate_callout_id()
+            bt = self.bt
+            client = args[0]
+            request = args[1]
+            if isinstance(request, str):
+                url = request
+            else:
+                url = request.url
+
+            parsed_url = urlparse(url)
+            host = parsed_url.hostname
+            port = parsed_url.port or ('443' if parsed_url.scheme == 'https' else '80')
+            url = parsed_url.path
+            host_port = "NA|" + host + "|" + str(port) + "|" + url
+            mName = fetch.__module__ + "." + fetch.__qualname__
+            headers_value = self.agent.get_nd_header(self.bt, host_port, unique_id).decode()
+
+            headers = None
+
+            if isinstance(request, str):
+                try:
+                    from tornado.httpclient import HTTPRequest
+                    req = HTTPRequest(request, **kwargs)
+                    headers = req._headers
+                except ImportError:
+                    pass
+            else:
+                headers = request._headers
+
+            if isinstance(headers, dict):
+                headers["CavNDFPInstance"] = headers_value
+            else:
+                try:
+                    from tornado.httputil import HTTPHeaders
+                    if isinstance(headers, HTTPHeaders):
+                        headers.add("CavNDFPInstance", headers_value)
+                except ImportError:
+                    pass
+
+            start_time = get_current_timestamp_in_us()
+            exit_call = self.http_call_begin(bt, host_port, url, mName, start_time, host_port, unique_id)
+        except Exception as e:
+            self.agent.logger.exception("Exception in HTTP Begin for Tornado Client {}".format(e))
+
+        future = fetch(*args, **kwargs)
+        future.add_done_callback(functools.partial(self.end, exit_call, mName, host_port, start_time, unique_id))
+
+        return future
+
+    def end(self, exit_call, mName, host_port, start_time, unique_id, future):
+        try:
+            if future._result:
+                current_time = get_current_timestamp_in_us()
+                duration = get_validated_duration(start_time, current_time, "Tornado Async HTTP Client")
+                status = future._result.code
+                self.http_call_end(self.bt, exit_call, mName, status, duration, start_time, host_port, unique_id)
+        except Exception as e:
+            self.agent.logger.exception("Exception in HTTP End for Tornado Client {}".format(e))
+
+
+def intercept_tornado_httpclient(agent, mod):
+    agent.logger.warning("Instrument module: tornado.httpclient{}".format(", mod: {}".format(mod) if mod else ""))
+    return TornadoAsyncHTTPClientInterceptor(agent, mod.AsyncHTTPClient).attach('fetch')

pythonagent/agent/probes/http/urllib3.py

@@ -0,0 +1,16 @@
+
+"""Intercept urllib3 to ensure that HTTPS is reported correctly.
+
+"""
+
+from __future__ import unicode_literals
+
+from . import HTTPConnectionInterceptor
+
+
+def intercept_urllib3(agent, mod):
+    agent.logger.warning("Instrument module: urllib3{}".format(", mod: {}".format(mod) if mod else ""))
+    # urllib3 1.8+ provides its own HTTPSConnection class.
+    if hasattr(mod, 'connection'):
+        #print("intercept_urllib3")
+        HTTPConnectionInterceptor.https_connection_classes.add(mod.connection.HTTPSConnection)

pythonagent/agent/probes/langchain/__init__.py

@@ -0,0 +1,21 @@
+#from .langchain_core import intercept_langchain_core_language_models
+#from .langchain_core import intercept_langchain_core_runnables
+from .langchain_core import intercept_langchain_core_tools
+
+from .langchain_community import intercept_langchain_community_vectorstores
+from .langchain_community import intercept_langchain_community_embeddings
+
+from .langchain_openai import intercept_langchain_openai
+
+#from .langgraph import intercept_langgraph_pregel, intercept_langgraph_utils_runnable
+
+__all__ = [#'intercept_langchain_core_language_models',
+           #'intercept_langchain_core_runnables',
+           'intercept_langchain_core_tools',
+           'intercept_langchain_community_vectorstores',
+           'intercept_langchain_community_embeddings',
+           'intercept_langchain_openai',
+           #'intercept_langgraph_pregel',
+           #'intercept_langgraph_utils_runnable'
+           ]
+

pythonagent/agent/probes/langchain/base_tool.py

@@ -0,0 +1,95 @@
+from ..base import ExitCallInterceptor
+from ..span import ToolSpan
+from pythonagent.utils import  get_current_timestamp_in_us, get_validated_duration, split_fp_topo
+import logging
+
+"""
+A wrapper class to instrument LangChain Tool 
+executions triggered by LLM outputs.
+Intercepts _invoke method to capture tool name
+input and output of the prompt.
+"""
+
+class LangChainBaseToolInterceptor(ExitCallInterceptor):
+    def _invoke(self, invoke, *args, **kwargs):
+
+        start_time, exit_call, bt, context, fqm, host_name_str = 0, 0, None, None, None, None
+        test_run, fpi, entry_fpi, nd_session_id, nv_session_id, page_id, full_fp = None, None, None, None, None, None, None
+
+        try:
+            tool_object, tool_parameters = args[0], args[1]
+            tool_name = tool_object.name
+
+            bt, callout_type, db = self.bt, "TOOLS", tool_name
+            host, port, username, query, query_params = "NA", "NA", "NA", "NA", None
+
+            host_name_str = f"NA|{host}|{port}|NA|{callout_type}|{db}|NA|NA|NA|{username}|NA"
+            fqm = f"{invoke.__module__}.{invoke.__qualname__}"
+
+            context = self.agent.get_transaction_context()
+            if context:
+                context.host_str = host_name_str
+                context.entry_point_fqm = fqm
+
+
+            start_time = get_current_timestamp_in_us()
+            if self.agent.logger.isEnabledFor(logging.INFO):
+                self.agent.logger.info("langchain base tool db call begin bt {} host {} query {} query_params {} start_time {}".format(bt, host, query, query_params, start_time))
+            exit_call = self.db_call_begin(bt, host, query, query_params, start_time)
+
+        except Exception as e:
+            self.agent.logger.exception("langchain base tool db call begin {}".format(e))
+
+        try:
+            result = invoke(*args, **kwargs)
+        except Exception as e:
+            self.agent.logger.exception("Error in application: {}".format(e))
+            try:
+                end_time = get_current_timestamp_in_us()
+                duration = get_validated_duration(start_time, end_time, "Langchain Tool")
+                if context:
+                    context.status_code = 500
+                    context.db_callout_time = 0
+                if self.agent.logger.isEnabledFor(logging.INFO):
+                    self.agent.logger.info("langchain base tool db call end bt {} exit_call {} duration {} start_time {} fqm {}host_name_str {}".format(bt, exit_call, duration, start_time, fqm, host_name_str))
+                self.db_call_end(bt, exit_call, duration, start_time, fqm, host_name_str)
+            except Exception as e:
+                self.agent.logger.exception("Exception in Original Method DB Call End for BaseTool invoke {}".format(e))
+
+            raise
+
+        try:
+            end_time = get_current_timestamp_in_us()
+            duration = get_validated_duration(start_time, end_time, "Langchain Tool")
+            if context:
+                context.status_code = 200
+                context.db_callout_time = 0
+
+            if self.agent.logger.isEnabledFor(logging.INFO):
+                self.agent.logger.info("langchain base tool db call end bt {} exit_call {} duration {} start_time {} fqm {}host_name_str {}".format(bt, exit_call, duration, start_time, fqm, host_name_str))
+            self.db_call_end(bt, exit_call, duration, start_time, fqm, host_name_str)
+
+            if self.agent.nf_enabled:
+                # Capturing input prompt and output response, formatting it, and sending to NF server
+                meta = { "span_kind": "tool",
+                         "input": str(args),
+                         "output": str(result)}
+
+                if context:
+                    fp_topo = context.fp_topo
+                    test_run, fpi, entry_fpi, nd_session_id, nv_session_id, page_id, full_fp = split_fp_topo(fp_topo)
+
+                tool_span = ToolSpan(fqm, start_time, duration, "ok", meta, test_run, fpi, entry_fpi, nd_session_id, nv_session_id, page_id, full_fp)
+                tool_span_json = tool_span.create_span_json()
+
+                if self.agent.logger.isEnabledFor(logging.INFO):
+                    self.agent.logger.info("langchain base tool sending tool span {}".format(tool_span_json))
+
+                # SDK API call to send logs to NF
+                self.agent.dump_log_message(bt, tool_span_json)
+
+        except Exception as e:
+            self.agent.logger.exception("langchain base tool db call begin {}".format(e))
+
+        return result
+