From c78eb44a0e458c49cf6386c3fad4b12e4e999f4e Mon Sep 17 00:00:00 2001 From: Anthony Mahanna Date: Fri, 3 Nov 2023 13:59:43 -0400 Subject: [PATCH 01/11] initial commit --- arango/client.py | 16 ++++++++++ arango/connection.py | 69 ++++++++++++++++++++++++++++++++------------ arango/exceptions.py | 8 +++++ docs/auth.rst | 5 ++++ tests/test_auth.py | 17 +++++++++-- 5 files changed, 94 insertions(+), 21 deletions(-) diff --git a/arango/client.py b/arango/client.py index a6bcab6d..11f9dfbd 100644 --- a/arango/client.py +++ b/arango/client.py @@ -171,6 +171,7 @@ def db( password: str = "", verify: bool = False, auth_method: str = "basic", + user_token: Optional[str] = None, superuser_token: Optional[str] = None, verify_certificate: bool = True, ) -> StandardDatabase: @@ -189,6 +190,10 @@ def db( refreshed automatically using ArangoDB username and password. This assumes that the clocks of the server and client are synchronized. :type auth_method: str + :param user_token: User generated token for user access. + If set, parameters **username**, **password** and **auth_method** + are ignored. This token is not refreshed automatically. + :type user_token: str :param superuser_token: User generated token for superuser access. If set, parameters **username**, **password** and **auth_method** are ignored. This token is not refreshed automatically. @@ -213,6 +218,17 @@ def db( deserializer=self._deserializer, superuser_token=superuser_token, ) + elif user_token is not None: + connection = JwtConnection( + hosts=self._hosts, + host_resolver=self._host_resolver, + sessions=self._sessions, + db_name=name, + http_client=self._http, + serializer=self._serializer, + deserializer=self._deserializer, + user_token=user_token, + ) elif auth_method.lower() == "basic": connection = BasicConnection( hosts=self._hosts, diff --git a/arango/connection.py b/arango/connection.py index 49aa7b67..cb996222 100644 --- a/arango/connection.py +++ b/arango/connection.py @@ -13,10 +13,16 @@ from typing import Any, Callable, Optional, Sequence, Set, Tuple, Union import jwt +from jwt.exceptions import ExpiredSignatureError from requests import ConnectionError, Session from requests_toolbelt import MultipartEncoder -from arango.exceptions import JWTAuthError, ServerConnectionError +from arango.exceptions import ( + JWTAuthError, + JWTExpiredError, + JWTRefreshError, + ServerConnectionError, +) from arango.http import HTTPClient from arango.request import Request from arango.resolver import HostResolver @@ -300,11 +306,12 @@ def __init__( host_resolver: HostResolver, sessions: Sequence[Session], db_name: str, - username: str, - password: str, http_client: HTTPClient, serializer: Callable[..., str], deserializer: Callable[[str], Any], + username: Optional[str] = None, + password: Optional[str] = None, + user_token: Optional[str] = None, ) -> None: super().__init__( hosts, @@ -323,7 +330,13 @@ def __init__( self._token: Optional[str] = None self._token_exp: int = sys.maxsize - self.refresh_token() + if user_token is not None: + self.set_token(user_token) + elif username is not None and password is not None: + self.refresh_token() + else: + m = "Either **user_token** or **username** & **password** must be set" + raise ValueError(m) def send_request(self, request: Request) -> Response: """Send an HTTP request to ArangoDB server. @@ -360,7 +373,12 @@ def refresh_token(self) -> None: :return: JWT token. :rtype: str + :raise arango.exceptions.JWTRefreshError: If missing username & password. + :raise arango.exceptions.JWTAuthError: If token retrieval fails. """ + if self._username is None or self._password is None: + raise JWTRefreshError("username and password must be set") + request = Request( method="post", endpoint="/_open/auth", @@ -374,21 +392,34 @@ def refresh_token(self) -> None: if not resp.is_success: raise JWTAuthError(resp, request) - self._token = resp.body["jwt"] - assert self._token is not None - - jwt_payload = jwt.decode( - self._token, - issuer="arangodb", - algorithms=["HS256"], - options={ - "require_exp": True, - "require_iat": True, - "verify_iat": True, - "verify_exp": True, - "verify_signature": False, - }, - ) + self.set_token(resp.body["jwt"]) + + def set_token(self, token: str) -> None: + """Set the JWT token. + + :param token: JWT token. + :type token: str + :raise arango.exceptions.JWTExpiredError: If the token is expired. + """ + assert token is not None + + try: + jwt_payload = jwt.decode( + token, + issuer="arangodb", + algorithms=["HS256"], + options={ + "require_exp": True, + "require_iat": True, + "verify_iat": True, + "verify_exp": True, + "verify_signature": False, + }, + ) + except ExpiredSignatureError: + raise JWTExpiredError("JWT token is expired") + + self._token = token self._token_exp = jwt_payload["exp"] self._auth_header = f"bearer {self._token}" diff --git a/arango/exceptions.py b/arango/exceptions.py index 8998f6c5..fb11f8d5 100644 --- a/arango/exceptions.py +++ b/arango/exceptions.py @@ -1014,3 +1014,11 @@ class JWTSecretListError(ArangoServerError): class JWTSecretReloadError(ArangoServerError): """Failed to reload JWT secrets.""" + + +class JWTRefreshError(ArangoClientError): + """Failed to refresh JWT token.""" + + +class JWTExpiredError(ArangoClientError): + """JWT token has expired.""" diff --git a/docs/auth.rst b/docs/auth.rst index a0cd9ac6..45c7171b 100644 --- a/docs/auth.rst +++ b/docs/auth.rst @@ -40,6 +40,7 @@ to work correctly. .. testcode:: from arango import ArangoClient + import os # Initialize the ArangoDB client. client = ArangoClient() @@ -59,6 +60,10 @@ to work correctly. # compensate for out-of-sync clocks between the client and server. db.conn.ext_leeway = 2 + # It is also possible to connect via a pre-generated JWT. + token = os.environ["ARANGODB_USER_JWT"] + db = client.db('test', user_token=token) + User generated JWT token can be used for superuser access. **Example:** diff --git a/tests/test_auth.py b/tests/test_auth.py index 9688799a..65c1dd52 100644 --- a/tests/test_auth.py +++ b/tests/test_auth.py @@ -2,6 +2,7 @@ from arango.errno import FORBIDDEN, HTTP_UNAUTHORIZED from arango.exceptions import ( JWTAuthError, + JWTExpiredError, JWTSecretListError, JWTSecretReloadError, ServerEncryptionError, @@ -37,7 +38,8 @@ def test_auth_basic(client, db_name, username, password): assert isinstance(db.properties(), dict) -def test_auth_jwt(client, db_name, username, password): +def test_auth_jwt(client, db_name, username, password, secret): + # Test JWT authentication with username and password. db = client.db( name=db_name, username=username, @@ -54,6 +56,13 @@ def test_auth_jwt(client, db_name, username, password): client.db(db_name, username, bad_password, auth_method="jwt") assert err.value.error_code == HTTP_UNAUTHORIZED + # Test JWT authentication with user token. + token = generate_jwt(secret) + db = client.db("_system", user_token=token) + assert isinstance(db.conn, JwtConnection) + assert isinstance(db.version(), str) + assert isinstance(db.properties(), dict) + # TODO re-examine commented out code def test_auth_superuser_token(client, db_name, root_password, secret): @@ -121,8 +130,12 @@ def test_auth_jwt_expiry(client, db_name, root_password, secret): db.conn._auth_header = f"bearer {expired_token}" assert isinstance(db.version(), str) - # Test correct error on token expiry. + # Test correct error on token expiry (superuser). db = client.db("_system", superuser_token=expired_token) with assert_raises(ServerVersionError) as err: db.version() assert err.value.error_code == FORBIDDEN + + # Test correct error on token expiry (user). + with assert_raises(JWTExpiredError) as err: + db = client.db("_system", user_token=expired_token) From 34aaa138e62dace4469c3da6874adba4625f3ef0 Mon Sep 17 00:00:00 2001 From: Anthony Mahanna Date: Fri, 3 Nov 2023 16:17:28 -0400 Subject: [PATCH 02/11] fix rst --- docs/auth.rst | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/docs/auth.rst b/docs/auth.rst index 45c7171b..31281f5c 100644 --- a/docs/auth.rst +++ b/docs/auth.rst @@ -40,7 +40,6 @@ to work correctly. .. testcode:: from arango import ArangoClient - import os # Initialize the ArangoDB client. client = ArangoClient() @@ -60,11 +59,7 @@ to work correctly. # compensate for out-of-sync clocks between the client and server. db.conn.ext_leeway = 2 - # It is also possible to connect via a pre-generated JWT. - token = os.environ["ARANGODB_USER_JWT"] - db = client.db('test', user_token=token) - -User generated JWT token can be used for superuser access. +User generated JWT token can be used for user and superuser access. **Example:** @@ -94,3 +89,9 @@ User generated JWT token can be used for superuser access. # Connect to "test" database as superuser using the token. db = client.db('test', superuser_token=token) + + # Connect to "test" database as user using the token. + db = client.db('test', user_token=token) + + # Manually set the token (JwtConnection only). + db.conn.set_token('new token') From 74e1b0bb4b6f826577094d9fb5b0cf186cfb4b86 Mon Sep 17 00:00:00 2001 From: Anthony Mahanna Date: Mon, 6 Nov 2023 11:21:48 -0500 Subject: [PATCH 03/11] `set_token` test --- tests/test_auth.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/test_auth.py b/tests/test_auth.py index 65c1dd52..91c7b870 100644 --- a/tests/test_auth.py +++ b/tests/test_auth.py @@ -139,3 +139,8 @@ def test_auth_jwt_expiry(client, db_name, root_password, secret): # Test correct error on token expiry (user). with assert_raises(JWTExpiredError) as err: db = client.db("_system", user_token=expired_token) + + # Test set_token() with expired token. + db = client.db("_system", user_token=generate_jwt(secret)) + with assert_raises(JWTExpiredError) as err: + db.conn.set_token(expired_token) From 7315740c31b4343bea394b42b5f226f623947aab Mon Sep 17 00:00:00 2001 From: Anthony Mahanna Date: Mon, 6 Nov 2023 12:07:25 -0500 Subject: [PATCH 04/11] update `user_token` docstring --- arango/client.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/arango/client.py b/arango/client.py index 11f9dfbd..74f8fbc0 100644 --- a/arango/client.py +++ b/arango/client.py @@ -192,7 +192,9 @@ def db( :type auth_method: str :param user_token: User generated token for user access. If set, parameters **username**, **password** and **auth_method** - are ignored. This token is not refreshed automatically. + are ignored. This token is not refreshed automatically. If automatic + token refresh is required, consider setting **auth_method** to "jwt" + and using the **username** and **password** parameters instead. :type user_token: str :param superuser_token: User generated token for superuser access. If set, parameters **username**, **password** and **auth_method** From f60ca8104028ab4d63afca37be7e163431f29423 Mon Sep 17 00:00:00 2001 From: Anthony Mahanna Date: Mon, 6 Nov 2023 12:07:41 -0500 Subject: [PATCH 05/11] new: `set_token` for `JwtSuperuserConnection` --- arango/connection.py | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/arango/connection.py b/arango/connection.py index cb996222..2fabef2c 100644 --- a/arango/connection.py +++ b/arango/connection.py @@ -475,3 +475,30 @@ def send_request(self, request: Request) -> Response: request.headers["Authorization"] = self._auth_header return self.process_request(host_index, request) + + def set_token(self, token: str) -> None: + """Set the JWT token. + + :param token: JWT token. + :type token: str + :raise arango.exceptions.JWTExpiredError: If the token is expired. + """ + assert token is not None + + try: + jwt.decode( + token, + issuer="arangodb", + algorithms=["HS256"], + options={ + "require_exp": True, + "require_iat": True, + "verify_iat": True, + "verify_exp": True, + "verify_signature": False, + }, + ) + except ExpiredSignatureError: + raise JWTExpiredError("JWT token is expired") + + self._auth_header = f"bearer {self._token}" From fa5739b4ba7140f9f16f470a2d515feb4e4bc19f Mon Sep 17 00:00:00 2001 From: Anthony Mahanna Date: Mon, 6 Nov 2023 12:07:52 -0500 Subject: [PATCH 06/11] add `set_token` example in rst --- docs/auth.rst | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/docs/auth.rst b/docs/auth.rst index 31281f5c..11f62985 100644 --- a/docs/auth.rst +++ b/docs/auth.rst @@ -93,5 +93,25 @@ User generated JWT token can be used for user and superuser access. # Connect to "test" database as user using the token. db = client.db('test', user_token=token) - # Manually set the token (JwtConnection only). +User and superuser tokens can be set on the connection object as well. + +**Example:** + +.. code-block:: python + + from arango import ArangoClient + + # Initialize the ArangoDB client. + client = ArangoClient() + + # Connect to "test" database as superuser using the token. + db = client.db('test', user_token='token') + + # Set the user token on the connection object. db.conn.set_token('new token') + + # Connect to "test" database as superuser using the token. + db = client.db('test', superuser_token='superuser token') + + # Set the user token on the connection object. + db.conn.set_token('new superuser token') From ededc850cc21368e35c7127209fe7f23863814bd Mon Sep 17 00:00:00 2001 From: Anthony Mahanna Date: Mon, 6 Nov 2023 12:08:22 -0500 Subject: [PATCH 07/11] fix: set_token --- arango/connection.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/arango/connection.py b/arango/connection.py index 2fabef2c..ae076268 100644 --- a/arango/connection.py +++ b/arango/connection.py @@ -501,4 +501,4 @@ def set_token(self, token: str) -> None: except ExpiredSignatureError: raise JWTExpiredError("JWT token is expired") - self._auth_header = f"bearer {self._token}" + self._auth_header = f"bearer {token}" From 5ab893d5924c13464e1615a79b6d7962fd96df76 Mon Sep 17 00:00:00 2001 From: Anthony Mahanna Date: Tue, 7 Nov 2023 11:06:20 -0500 Subject: [PATCH 08/11] update docstring --- arango/client.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/arango/client.py b/arango/client.py index 74f8fbc0..13d53cb1 100644 --- a/arango/client.py +++ b/arango/client.py @@ -194,11 +194,13 @@ def db( If set, parameters **username**, **password** and **auth_method** are ignored. This token is not refreshed automatically. If automatic token refresh is required, consider setting **auth_method** to "jwt" - and using the **username** and **password** parameters instead. + and using the **username** and **password** parameters instead. Token + expiry will be checked. :type user_token: str :param superuser_token: User generated token for superuser access. If set, parameters **username**, **password** and **auth_method** - are ignored. This token is not refreshed automatically. + are ignored. This token is not refreshed automatically. Token + expiry will not be checked. :type superuser_token: str :param verify_certificate: Verify TLS certificates. :type verify_certificate: bool From 85e14a953c1f22c9288fe5b223a91ebcb688f2cd Mon Sep 17 00:00:00 2001 From: Anthony Mahanna Date: Tue, 7 Nov 2023 11:06:30 -0500 Subject: [PATCH 09/11] update error msg --- arango/connection.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/arango/connection.py b/arango/connection.py index ae076268..3daa4585 100644 --- a/arango/connection.py +++ b/arango/connection.py @@ -209,7 +209,7 @@ def ping(self) -> int: request = Request(method="get", endpoint="/_api/collection") resp = self.send_request(request) if resp.status_code in {401, 403}: - raise ServerConnectionError("bad username and/or password") + raise ServerConnectionError("bad username/password or token is expired") if not resp.is_success: # pragma: no cover raise ServerConnectionError(resp.error_message or "bad server response") return resp.status_code From 08f77e9a2ace57c04432d6bb921e60a11de0596b Mon Sep 17 00:00:00 2001 From: Anthony Mahanna Date: Tue, 7 Nov 2023 11:06:39 -0500 Subject: [PATCH 10/11] update `test_auth_jwt_expiry` --- tests/test_auth.py | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/tests/test_auth.py b/tests/test_auth.py index 91c7b870..8b68fd36 100644 --- a/tests/test_auth.py +++ b/tests/test_auth.py @@ -9,6 +9,7 @@ ServerTLSError, ServerTLSReloadError, ServerVersionError, + ServerConnectionError, ) from tests.helpers import assert_raises, generate_jwt, generate_string @@ -125,22 +126,32 @@ def test_auth_superuser_token(client, db_name, root_password, secret): def test_auth_jwt_expiry(client, db_name, root_password, secret): # Test automatic token refresh on expired token. db = client.db("_system", "root", root_password, auth_method="jwt") + valid_token = generate_jwt(secret) expired_token = generate_jwt(secret, exp=-1000) db.conn._token = expired_token db.conn._auth_header = f"bearer {expired_token}" assert isinstance(db.version(), str) - # Test correct error on token expiry (superuser). + # Test expiry error on db instantiation (superuser) + with assert_raises(ServerConnectionError) as err: + client.db("_system", superuser_token=expired_token, verify=True) + + # Test expiry error on db version (superuser) db = client.db("_system", superuser_token=expired_token) with assert_raises(ServerVersionError) as err: db.version() assert err.value.error_code == FORBIDDEN - # Test correct error on token expiry (user). + # Test expiry error on set_token (superuser). + db = client.db("_system", superuser_token=valid_token) + with assert_raises(JWTExpiredError) as err: + db.conn.set_token(expired_token) + + # Test expiry error on db instantiation (user) with assert_raises(JWTExpiredError) as err: db = client.db("_system", user_token=expired_token) - # Test set_token() with expired token. - db = client.db("_system", user_token=generate_jwt(secret)) + # Test expiry error on set_token (user). + db = client.db("_system", user_token=valid_token) with assert_raises(JWTExpiredError) as err: db.conn.set_token(expired_token) From 810e3f1c3d801cab37e48ce02bc2a564534b0de7 Mon Sep 17 00:00:00 2001 From: Anthony Mahanna Date: Tue, 7 Nov 2023 11:08:31 -0500 Subject: [PATCH 11/11] fix isort --- tests/test_auth.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_auth.py b/tests/test_auth.py index 8b68fd36..0f747563 100644 --- a/tests/test_auth.py +++ b/tests/test_auth.py @@ -5,11 +5,11 @@ JWTExpiredError, JWTSecretListError, JWTSecretReloadError, + ServerConnectionError, ServerEncryptionError, ServerTLSError, ServerTLSReloadError, ServerVersionError, - ServerConnectionError, ) from tests.helpers import assert_raises, generate_jwt, generate_string