From e9f1df667faa1c412a82bca576c80d007a8e788f Mon Sep 17 00:00:00 2001
From: Steffen Klemer <steffen.klemer@gwdg.de>
Date: Sun, 26 Nov 2023 17:59:10 +0100
Subject: [PATCH] Add support for generic auth headers

Add a auth_header attribute that is used verbatim
as Authorization header (i.e. without prefixing `Token `).
We use this to support oauth2 flows with Bearer tokens
(patch for nautobot will come soon).
---
 pynautobot/core/api.py      | 19 +++++++++++++++++--
 pynautobot/core/app.py      | 10 +++++-----
 pynautobot/core/endpoint.py | 20 ++++++++++----------
 pynautobot/core/query.py    | 13 +++++++++----
 pynautobot/core/response.py |  6 +++---
 pynautobot/models/dcim.py   |  2 +-
 6 files changed, 45 insertions(+), 25 deletions(-)

diff --git a/pynautobot/core/api.py b/pynautobot/core/api.py
index 0cbc49c..462b880 100644
--- a/pynautobot/core/api.py
+++ b/pynautobot/core/api.py
@@ -53,6 +53,7 @@ class Api(object):
     :param str url: The base URL to the instance of Nautobot you
         wish to connect to.
     :param str token: Your Nautobot token.
+    :param str auth_header: Content of the ``Authorization``-Header, e.g. for Bearer Token.
     :param bool,optional threading: Set to True to use threading in ``.all()``
         and ``.filter()`` requests.
     :param int,optional max_workers: Set the maximum workers for threading in ``.all()``
@@ -71,12 +72,22 @@ class Api(object):
     ...     token='d6f4e314a5b5fefd164995169f28ae32d987704f'
     ... )
     >>> nb.dcim.devices.all()
+
+    or for oauth2 flows
+
+    >>> import pynautobot
+    >>> nb = pynautobot.api(
+    ...     'http://localhost:8000',
+    ...     auth_header='Bearer keycloak eyJhbGciOiJSUzI1NiIsInR5cCIgOiAiSl...'
+    ... )
+    >>> nb.dcim.devices.all()
     """
 
     def __init__(
         self,
         url,
         token=None,
+        auth_header=None,
         threading=False,
         max_workers=4,
         api_version=None,
@@ -85,7 +96,11 @@ def __init__(
     ):
         base_url = "{}/api".format(url if url[-1] != "/" else url[:-1])
         self.token = token
-        self.headers = {"Authorization": f"Token {self.token}"}
+        if auth_header:
+            self.auth_header = auth_header
+        else:
+            self.auth_header = f"Token {self.token}"
+        self.headers = {"Authorization": self.auth_header}
         self.base_url = base_url
         self.http_session = requests.Session()
         self.http_session.verify = verify
@@ -198,7 +213,7 @@ def status(self):
         """
         status = Request(
             base=self.base_url,
-            token=self.token,
+            auth_header=self.auth_header,
             http_session=self.http_session,
             api_version=self.api_version,
         ).get_status()
diff --git a/pynautobot/core/app.py b/pynautobot/core/app.py
index 8d04e4e..1168cfe 100644
--- a/pynautobot/core/app.py
+++ b/pynautobot/core/app.py
@@ -80,7 +80,7 @@ def choices(self):
 
         self._choices = Request(
             base="{}/{}/_choices/".format(self.api.base_url, self.name),
-            token=self.api.token,
+            auth_header=self.api.auth_header,
             http_session=self.api.http_session,
         ).get()
 
@@ -144,7 +144,7 @@ def get_custom_fields(self):
         """
         return Request(
             base=f"{self.api.base_url}/{self.name}/custom-fields/",
-            token=self.api.token,
+            auth_header=self.api.auth_header,
             http_session=self.api.http_session,
         ).get()
 
@@ -176,7 +176,7 @@ def get_custom_field_choices(self):
         """
         return Request(
             base=f"{self.api.base_url}/{self.name}/custom-field-choices/",
-            token=self.api.token,
+            auth_header=self.api.auth_header,
             http_session=self.api.http_session,
         ).get()
 
@@ -201,7 +201,7 @@ def config(self):
                 self.api.base_url,
                 self.name,
             ),
-            token=self.api.token,
+            auth_header=self.api.auth_header,
             http_session=self.api.http_session,
         ).get()
         return config
@@ -241,7 +241,7 @@ def installed_plugins(self):
             base="{}/plugins/installed-plugins".format(
                 self.api.base_url,
             ),
-            token=self.api.token,
+            auth_header=self.api.auth_header,
             http_session=self.api.http_session,
         ).get()
         return installed_plugins
diff --git a/pynautobot/core/endpoint.py b/pynautobot/core/endpoint.py
index c2dc4f6..bd6a288 100644
--- a/pynautobot/core/endpoint.py
+++ b/pynautobot/core/endpoint.py
@@ -54,7 +54,7 @@ def __init__(self, api, app, name, model=None):
         self.name = name.replace("_", "-")
         self.api = api
         self.base_url = api.base_url
-        self.token = api.token
+        self.auth_header = api.auth_header
         self.url = "{base_url}/{app}/{endpoint}".format(
             base_url=self.base_url,
             app=app.name,
@@ -101,7 +101,7 @@ def all(self, api_version=None):
         api_version = api_version or self.api.api_version
         req = Request(
             base="{}/".format(self.url),
-            token=self.token,
+            auth_header=self.auth_header,
             http_session=self.api.http_session,
             threading=self.api.threading,
             max_workers=self.api.max_workers,
@@ -166,7 +166,7 @@ def get(self, *args, **kwargs):
         req = Request(
             key=key,
             base=self.url,
-            token=self.token,
+            auth_header=self.auth_header,
             http_session=self.api.http_session,
             api_version=api_version,
         )
@@ -237,7 +237,7 @@ def filter(self, *args, api_version=None, **kwargs):
         req = Request(
             filters=kwargs,
             base=self.url,
-            token=self.token,
+            auth_header=self.auth_header,
             http_session=self.api.http_session,
             threading=self.api.threading,
             api_version=api_version,
@@ -302,7 +302,7 @@ def create(self, *args, api_version=None, **kwargs):
 
         req = Request(
             base=self.url,
-            token=self.token,
+            auth_header=self.auth_header,
             http_session=self.api.http_session,
             api_version=api_version,
         ).post(args[0] if args else kwargs)
@@ -338,7 +338,7 @@ def update(self, id: str, data: Dict[str, any]):
         req = Request(
             key=id,
             base=self.url,
-            token=self.api.token,
+            auth_header=self.api.auth_header,
             http_session=self.api.http_session,
             api_version=self.api.api_version,
         )
@@ -385,7 +385,7 @@ def choices(self, api_version=None):
 
         req = Request(
             base=self.url,
-            token=self.api.token,
+            auth_header=self.api.auth_header,
             http_session=self.api.http_session,
             api_version=api_version,
         ).options()
@@ -442,7 +442,7 @@ def count(self, *args, api_version=None, **kwargs):
         api_version = api_version or self.api.api_version
 
         ret = Request(
-            filters=kwargs, base=self.url, token=self.token, http_session=self.api.http_session, api_version=api_version
+            filters=kwargs, base=self.url, auth_header=self.auth_header, http_session=self.api.http_session, api_version=api_version
         )
 
         return ret.get_count()
@@ -462,7 +462,7 @@ def __init__(self, parent_obj, name, custom_return=None):
 
         self.request_kwargs = dict(
             base=self.url,
-            token=parent_obj.api.token,
+            auth_header=parent_obj.api.auth_header,
             http_session=parent_obj.api.http_session,
         )
 
@@ -560,7 +560,7 @@ def run(self, *args, api_version=None, **kwargs):
 
         req = Request(
             base=job_run_url,
-            token=self.token,
+            auth_header=self.auth_header,
             http_session=self.api.http_session,
             api_version=api_version,
         ).post(args[0] if args else kwargs)
diff --git a/pynautobot/core/query.py b/pynautobot/core/query.py
index 6b3cb46..b3038aa 100644
--- a/pynautobot/core/query.py
+++ b/pynautobot/core/query.py
@@ -132,6 +132,7 @@ def __init__(
         filters=None,
         key=None,
         token=None,
+        auth_header=None,
         threading=False,
         max_workers=4,
         api_version=None,
@@ -152,6 +153,10 @@ def __init__(
         self.filters = filters
         self.key = key
         self.token = token
+        if auth_header:
+            self.auth_header = auth_header
+        else:
+            self.auth_header = f"Token {self.token}"
         self.http_session = http_session
         self.url = self.base if not key else "{}{}/".format(self.base, key)
         self.threading = threading
@@ -214,8 +219,8 @@ def get_status(self):
         :Raises: RequestError if request is not successful.
         """
         headers = {"Content-Type": "application/json;"}
-        if self.token:
-            headers["authorization"] = "Token {}".format(self.token)
+        if self.auth_header:
+            headers["authorization"] = self.auth_header
 
         if self.api_version:
             headers["accept"] = f"application/json; version={self.api_version}"
@@ -246,8 +251,8 @@ def _make_call(self, verb="get", url_override=None, add_params=None, data=None):
         else:
             headers = {"accept": "application/json;"}
 
-        if self.token:
-            headers["authorization"] = "Token {}".format(self.token)
+        if self.auth_header:
+            headers["authorization"] = self.auth_header
 
         if self.api_version:
             headers["accept"] = f"application/json; version={self.api_version}"
diff --git a/pynautobot/core/response.py b/pynautobot/core/response.py
index d2874d6..e790353 100644
--- a/pynautobot/core/response.py
+++ b/pynautobot/core/response.py
@@ -298,7 +298,7 @@ def full_details(self):
         if self.url:
             req = Request(
                 base=self.url,
-                token=self.api.token,
+                auth_header=self.api.auth_header,
                 http_session=self.api.http_session,
                 api_version=self.api.api_version,
             )
@@ -385,7 +385,7 @@ def save(self):
                 req = Request(
                     key=self.id,
                     base=self.endpoint.url,
-                    token=self.api.token,
+                    auth_header=self.api.auth_header,
                     http_session=self.api.http_session,
                     api_version=self.api.api_version,
                 )
@@ -433,7 +433,7 @@ def delete(self):
         req = Request(
             key=self.id,
             base=self.endpoint.url,
-            token=self.api.token,
+            auth_header=self.api.auth_header,
             http_session=self.api.http_session,
             api_version=self.api.api_version,
         )
diff --git a/pynautobot/models/dcim.py b/pynautobot/models/dcim.py
index 89af558..926a424 100644
--- a/pynautobot/models/dcim.py
+++ b/pynautobot/models/dcim.py
@@ -28,7 +28,7 @@ def trace(self):
         req = Request(
             key=str(self.id) + "/trace",
             base=self.endpoint.url,
-            token=self.api.token,
+            auth_header=self.api.auth_header,
             http_session=self.api.http_session,
         )
         uri_to_obj_class_map = {