From a04024c62642fecf21f76ce469eaa720ed2080f0 Mon Sep 17 00:00:00 2001
From: Nigel Small <nigel.small@neotechnology.com>
Date: Fri, 2 Dec 2016 17:41:38 +0000
Subject: [PATCH] Custom auth (from Pontus' original work)

---
 neo4j/v1/session.py | 27 ++++++++++++++++++++++++---
 test/test_driver.py | 26 +++++++++++++++++++++++++-
 2 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/neo4j/v1/session.py b/neo4j/v1/session.py
index 95af0ab54..725454940 100644
--- a/neo4j/v1/session.py
+++ b/neo4j/v1/session.py
@@ -57,10 +57,17 @@ class AuthToken(object):
     """ Container for auth information
     """
 
-    def __init__(self, scheme, principal, credentials):
+    #: By default we should not send any realm
+    realm = None
+
+    def __init__(self, scheme, principal, credentials, realm=None, **parameters):
         self.scheme = scheme
         self.principal = principal
         self.credentials = credentials
+        if realm:
+            self.realm = realm
+        if parameters:
+            self.parameters = parameters
 
 
 class GraphDatabase(object):
@@ -665,14 +672,28 @@ def __ne__(self, other):
         return not self.__eq__(other)
 
 
-def basic_auth(user, password):
+def basic_auth(user, password, realm=None):
     """ Generate a basic auth token for a given user and password.
 
     :param user: user name
     :param password: current password
+    :param realm: specifies the authentication provider
+    :return: auth token for use with :meth:`GraphDatabase.driver`
+    """
+    return AuthToken("basic", user, password, realm)
+
+
+def custom_auth(principal, credentials, realm, scheme, **parameters):
+    """ Generate a basic auth token for a given user and password.
+
+    :param principal: specifies who is being authenticated
+    :param credentials: authenticates the principal
+    :param realm: specifies the authentication provider
+    :param scheme: specifies the type of authentication
+    :param parameters: parameters passed along to the authenticatin provider
     :return: auth token for use with :meth:`GraphDatabase.driver`
     """
-    return AuthToken("basic", user, password)
+    return AuthToken(scheme, principal, credentials, realm, **parameters)
 
 
 def parse_address(address):
diff --git a/test/test_driver.py b/test/test_driver.py
index de17503fc..93ff0d70b 100644
--- a/test/test_driver.py
+++ b/test/test_driver.py
@@ -25,7 +25,7 @@
 
 from neo4j.v1 import ServiceUnavailable, ProtocolError, READ_ACCESS, WRITE_ACCESS, \
     TRUST_ON_FIRST_USE, TRUST_CUSTOM_CA_SIGNED_CERTIFICATES, GraphDatabase, basic_auth, \
-    SSL_AVAILABLE, SessionExpired, DirectDriver
+    custom_auth, SSL_AVAILABLE, SessionExpired, DirectDriver
 from test.util import ServerTestCase
 
 BOLT_URI = "bolt://localhost:7687"
@@ -68,6 +68,30 @@ def test_fail_nicely_when_connecting_to_http_port(self):
         with self.assertRaises(ServiceUnavailable) as context:
             driver.session()
 
+    def test_can_provide_realm_with_basic_auth_token(self):
+        token = basic_auth("neotest", "neotest", "native")
+        driver = GraphDatabase.driver("bolt://localhost", auth=token)
+        session = driver.session()
+        result = session.run("RETURN 1").consume()
+        session.close()
+        assert result is not None
+
+    def test_can_create_custom_auth_token(self):
+        token = custom_auth("neotest", "neotest", "native", "basic")
+        driver = GraphDatabase.driver("bolt://localhost", auth=token)
+        session = driver.session()
+        result = session.run("RETURN 1").consume()
+        session.close()
+        assert result is not None
+
+    def test_can_create_custom_auth_token_with_additional_parameters(self):
+        token = custom_auth("neotest", "neotest", "native", "basic", secret=42)
+        driver = GraphDatabase.driver("bolt://localhost", auth=token)
+        session = driver.session()
+        result = session.run("RETURN 1").consume()
+        session.close()
+        assert result is not None
+
 
 class DirectDriverTestCase(ServerTestCase):