From 18358db8e879351c11ba868590fe3eec1f8de65e Mon Sep 17 00:00:00 2001
From: Anthony Ruhier <anthony.ruhier@gmail.com>
Date: Mon, 29 May 2017 15:35:25 +0200
Subject: [PATCH] Add token authentication

---
 README.md         | 14 ++++++++++++++
 netboxapi/api.py  | 24 +++++++++++++++++++++++-
 tests/test_api.py | 14 ++++++++++++++
 3 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 5d32574..bd570ae 100644
--- a/README.md
+++ b/README.md
@@ -21,9 +21,23 @@ from netboxapi import NetboxAPI
 Initialize a new `NetboxAPI` object:
 
 ```python
+netbox_api = NetboxAPI(url="netbox.example.com")
+
+# or if you enabled the authentication
 netbox_api = NetboxAPI(
     url="netbox.example.com", username="user", password="password"
 )
+
+# or if you have generated a token
+netbox_api = NetboxAPI(
+    url="netbox.example.com", token="token"
+)
+
+# but the following is useless, as the token will not be used
+netbox_api = NetboxAPI(
+    url="netbox.example.com", username="user", password="password",
+    token="token"
+)
 ```
 
 Then use multiple available methods to interact with the api:
diff --git a/netboxapi/api.py b/netboxapi/api.py
index f52fe4e..b280c52 100644
--- a/netboxapi/api.py
+++ b/netboxapi/api.py
@@ -3,10 +3,28 @@
 from urllib.parse import urljoin
 
 
+class _HTTPTokenAuth(requests.auth.AuthBase):
+    """HTTP Basic Authentication with token."""
+
+    def __init__(self, token):
+        self.token = token
+
+    def __eq__(self, other):
+        return self.token == getattr(other, 'token', None)
+
+    def __ne__(self, other):
+        return not self == other
+
+    def __call__(self, r):
+        r.headers["Authorization"] = "Token {}".format(self.token)
+        return r
+
+
 class NetboxAPI():
-    def __init__(self, url, username=None, password=None):
+    def __init__(self, url, username=None, password=None, token=None):
         self.username = username
         self.password = password
+        self.token = token
         self.url = url.rstrip('/')
 
         self.session = requests.Session()
@@ -43,6 +61,10 @@ def _generic_http_method_request(self, method, url, **kwargs):
             response = http_method(
                 url, auth=(self.username, self.password), **kwargs
             )
+        elif self.token:
+            response = http_method(
+                url, auth=_HTTPTokenAuth(self.token), **kwargs
+            )
         else:
             response = http_method(url, **kwargs)
 
diff --git a/tests/test_api.py b/tests/test_api.py
index 9b65d3c..e5d5d63 100644
--- a/tests/test_api.py
+++ b/tests/test_api.py
@@ -3,12 +3,14 @@
 import requests_mock
 
 from netboxapi import NetboxAPI
+from netboxapi.api import _HTTPTokenAuth
 
 
 class TestNetboxAPI():
     url = "http://localhost/"
     login = "login"
     password = "password"
+    token = "testing_token"
 
     @pytest.fixture()
     def prepared_api(self):
@@ -29,6 +31,10 @@ def test_get_loggedin(self, **kwargs):
         prepared_api = NetboxAPI(self.url, self.login, self.password)
         self._generic_test_http_method_request(prepared_api, "get")
 
+    def test_get_loggedin_token(self, **kwargs):
+        prepared_api = NetboxAPI(self.url, token=self.token)
+        self._generic_test_http_method_request(prepared_api, "get")
+
     def test_post(self, prepared_api, **kwargs):
         self._generic_test_http_method_request(prepared_api, "post")
 
@@ -51,3 +57,11 @@ def _generic_test_http_method_request(self, prepared_api, method):
             m.register_uri(method, url, json=expected_json)
             response = getattr(prepared_api, method)(url)
         assert response == expected_json
+
+
+class TestHTTPTokenAuth():
+    def test_eq(self):
+        assert _HTTPTokenAuth("test") == _HTTPTokenAuth("test")
+
+    def test_not_eq(self):
+        assert _HTTPTokenAuth("test") != _HTTPTokenAuth("test1")