From af135be2da221f388ef4b32dad352d63f0490678 Mon Sep 17 00:00:00 2001
From: Lukas Garberg <lukas@spritelink.net>
Date: Wed, 6 Mar 2024 08:22:54 +0100
Subject: [PATCH] pynipap: Fix TLS support after tracing feature

The implementation of OpenTelemetry tracing broke TLS support in
pynipap. This commit fixes that.
---
 pynipap/pynipap.py | 24 +++++++++++++++++-------
 pynipap/tracing.py | 35 ++++++++++++++++++++++++++++++++++-
 2 files changed, 51 insertions(+), 8 deletions(-)

diff --git a/pynipap/pynipap.py b/pynipap/pynipap.py
index 0a1c31eb2..d405b5c16 100644
--- a/pynipap/pynipap.py
+++ b/pynipap/pynipap.py
@@ -207,6 +207,9 @@
 import sys
 import logging
 import xmlrpc.client as xmlrpclib
+import urllib.parse
+
+from tracing import create_span
 
 __version__		= "0.32.3"
 __author__		= "Kristian Larsson, Lukas Garberg"
@@ -253,13 +256,6 @@ def __init__(self, options = None):
             self.options = options
 
 
-try:
-    from tracing import create_span, TracingXMLTransport
-    xml_transport = TracingXMLTransport
-except ImportError:
-    xml_transport = xmlrpclib.Transport
-
-
 class XMLRPCConnection:
     """ Handles a shared XML-RPC connection.
     """
@@ -281,6 +277,20 @@ def __init__(self):
         if xmlrpc_uri is None:
             raise NipapError('XML-RPC URI not specified')
 
+        p = urllib.parse.urlsplit(xmlrpc_uri)
+
+        try:
+            from tracing import TracingXMLTransport, TracingXMLSafeTransport
+            if p.scheme == "http":
+                xml_transport = TracingXMLTransport
+            elif p.scheme == "https":
+                xml_transport = TracingXMLSafeTransport
+        except ImportError:
+            if p.scheme == "http":
+                xml_transport = xmlrpclib.Transport
+            elif p.scheme == "https":
+                xml_transport = xmlrpclib.SafeTransport
+
         # creating new instance
         self.connection = xmlrpclib.ServerProxy(xmlrpc_uri,
                                                 transport=xml_transport(),
diff --git a/pynipap/tracing.py b/pynipap/tracing.py
index 0a779efaf..29bc3c336 100644
--- a/pynipap/tracing.py
+++ b/pynipap/tracing.py
@@ -1,7 +1,6 @@
 import inspect
 import sys
 from functools import wraps
-from pynipap import AuthOptions
 
 if sys.version_info[0] < 3:
     import xmlrpclib
@@ -45,6 +44,7 @@ def decorated(*args, **kwargs):
             """
             """
 
+            from pynipap import AuthOptions
             signature = inspect.signature(f)
 
             if args[0].__class__ == type:
@@ -103,6 +103,39 @@ def send_content(self, connection, request_body):
                     current_span.get_span_context().span_id)[2:].zfill(16) + "-01")
 
             super().send_content(connection, request_body)
+
+
+    class TracingXMLSafeTransport(xmlrpclib.SafeTransport):
+
+        def __init__(self, use_datetime=False, use_builtin_types=False,
+                     *, headers=(), context=None):
+            super().__init__(use_datetime=use_datetime,
+                             use_builtin_types=use_builtin_types,
+                             headers=headers, context=context)
+
+        def request(self, host, handler, request_body, verbose=False):
+            with tracer.start_as_current_span("POST XML request", context.get_current(), SpanKind.CLIENT):
+                current_span = trace.get_current_span()
+                try:
+                    result = super().request(host, handler, request_body, verbose)
+                    current_span.set_attribute("http.status_code", 200)
+                except xmlrpclib.ProtocolError as protocolError:
+                    current_span.set_attribute("http.status_code", protocolError.errcode)
+                    current_span.record_exception(protocolError)
+                    raise protocolError
+            return result
+
+        def send_content(self, connection, request_body):
+            current_span = trace.get_current_span()
+            if current_span != INVALID_SPAN:
+                current_span.set_attribute("net.peer.ip", connection.host)
+                current_span.set_attribute("net.peer.port", connection.port)
+                current_span.set_attribute("net.peer.transport", "ip_tcp")
+                connection.putheader("traceparent", "00-" + hex(current_span.get_span_context().trace_id)[2:].zfill(32) + "-" + hex(
+                    current_span.get_span_context().span_id)[2:].zfill(16) + "-01")
+
+            super().send_content(connection, request_body)
+
 except ImportError:
     def create_span(f):
         @wraps(f)