From de68ace20343e625d851c1d55d16805c7500e836 Mon Sep 17 00:00:00 2001 From: Noah Gorstein Date: Thu, 22 Jun 2023 09:07:13 -0400 Subject: [PATCH] [CLOUD-1366]: add support for impersonating users (#150) --- stardog/admin.py | 8 ++++++-- stardog/connection.py | 10 +++++++++- stardog/http/client.py | 4 ++++ test/test_admin_basic.py | 22 ++++++++++++++++++++++ test/test_connection.py | 24 +++++++++++++++++++++++- 5 files changed, 64 insertions(+), 4 deletions(-) diff --git a/stardog/admin.py b/stardog/admin.py index 281cc32..0863c01 100644 --- a/stardog/admin.py +++ b/stardog/admin.py @@ -26,6 +26,7 @@ def __init__( username: object = None, password: object = None, auth: object = None, + run_as: str = None, ) -> None: """Initializes an admin connection to a Stardog server. @@ -36,8 +37,9 @@ def __init__( Defaults to `admin` password (str, optional): Password to use in the connection. Defaults to `admin` - auth (requests.auth.AuthBase, optional): requests Authentication object. + auth (requests.auth.AuthBase, optional): requests Authentication object. Defaults to `None` + run_as (str, optional): the User to impersonate auth and username/password should not be used together. If the are the value of `auth` will take precedent. @@ -45,7 +47,9 @@ def __init__( >>> admin = Admin(endpoint='http://localhost:9999', username='admin', password='admin') """ - self.client = client.Client(endpoint, None, username, password, auth=auth) + self.client = client.Client( + endpoint, None, username, password, auth=auth, run_as=run_as + ) # ensure the server is alive and at the specified location self.alive() diff --git a/stardog/connection.py b/stardog/connection.py index 8f73c4b..522835b 100644 --- a/stardog/connection.py +++ b/stardog/connection.py @@ -25,6 +25,7 @@ def __init__( password=None, auth=None, session=None, + run_as=None, ): """Initializes a connection to a Stardog database. @@ -38,13 +39,20 @@ def __init__( Defaults to `None` session (requests.session.Session, optional): requests Session object. Defaults to `None` + run_as (str, optional): the user to impersonate Examples: >>> conn = Connection('db', endpoint='http://localhost:9999', username='admin', password='admin') """ self.client = client.Client( - endpoint, database, username, password, auth=auth, session=session + endpoint, + database, + username, + password, + auth=auth, + session=session, + run_as=run_as, ) self.transaction = None diff --git a/stardog/http/client.py b/stardog/http/client.py index 0162616..8ffcd2b 100644 --- a/stardog/http/client.py +++ b/stardog/http/client.py @@ -18,6 +18,7 @@ def __init__( password=None, session=None, auth=None, + run_as=None, ): self.url = endpoint if endpoint else self.DEFAULT_ENDPOINT @@ -45,6 +46,9 @@ def __init__( ) self.session.auth = auth + if run_as: + self.session.headers.update({"SD-Run-As": run_as}) + def post(self, path, **kwargs): return self.__wrap(self.session.post(self.url + path, **kwargs)) diff --git a/test/test_admin_basic.py b/test/test_admin_basic.py index b1a2b0a..cbde1d2 100644 --- a/test/test_admin_basic.py +++ b/test/test_admin_basic.py @@ -88,6 +88,28 @@ def music_options(self): } +class TestUserImpersonation(TestStardog): + def test_impersonating_user_databases_visibility(self, conn_string, user, db): + + with admin.Admin( + endpoint=conn_string["endpoint"], + ) as admin_user: + databases_admin_can_see = [db.name for db in admin_user.databases()] + with admin.Admin( + endpoint=conn_string["endpoint"], + run_as=user.name, + ) as admin_impersonating_user: + databases_impersonated_user_can_see = [ + db.name for db in admin_impersonating_user.databases() + ] + assert len(databases_impersonated_user_can_see) == 0 + + # for cluster tests in Circle, catalog is disabled so the exact number of dbs + # varies (2 for single node, 1 for cluster since catalog isn't created) + assert len(databases_admin_can_see) > 0 + assert db.name in databases_admin_can_see + + class TestUsers(TestStardog): def test_user_creation(self, admin, user): assert len(admin.users()) == len(default_users) + 1 diff --git a/test/test_connection.py b/test/test_connection.py index 3b64785..3d73d69 100644 --- a/test/test_connection.py +++ b/test/test_connection.py @@ -1,6 +1,6 @@ import pytest -from stardog import admin, connection, content, content_types, exceptions +from stardog import connection, content, content_types, exceptions def starwars_contents() -> list: @@ -100,6 +100,28 @@ def test_export(db, conn): assert b"default_obj" not in named_export +def test_user_impersonation(conn_string, db, user): + with connection.Connection( + endpoint=conn_string["endpoint"], database=db.name + ) as admin_regular_conn: + # add some data to query + admin_regular_conn.begin() + admin_regular_conn.clear() + admin_regular_conn.add(content.File("test/data/starwars.ttl")) + admin_regular_conn.commit() + + # confirm admin can read the data + q = admin_regular_conn.select('select * {?s :name "Luke Skywalker"}') + assert len(q["results"]["bindings"]) == 1 + + # attempting to query database as the user the admin is impersonating should return an unauthorized error + with connection.Connection( + endpoint=conn_string["endpoint"], database=db.name, run_as=user.name + ) as admin_impersonating_user: + with pytest.raises(exceptions.StardogException, match="403"): + admin_impersonating_user.select('select * {?s :name "Luke Skywalker"}') + + @pytest.mark.dbname("pystardog-test-database") @pytest.mark.conn_dbname("pystardog-test-database") def test_queries(db, conn):