Updates for multiregion changes

shilpabijam · shilpabijam · commit 70d0cb80f1b3 · 2021-03-12T17:17:46.000Z
diff --git a/example_sdkrc b/example_sdkrc
@@ -1,9 +1,11 @@
 # SPLUNK HOSTS
 export SPLUNK_AUTH_HOST=<REDACTED>
+export SPLUNK_AUTH_HOST_2=<REDACTED>
 export SPLUNK_AUTHN_URL=<REDACTED>
 export SPLUNK_HOST=<REDACTED>
 export SPLUNK_APP_SERVER=<REDACTED>
 export SPLUNK_DEBUG=<REDACTED>
+export SPLUNK_HOST_2=<REDACTED>
 
 # Choose one set of vars for PKCE or CLIENT auth flows
 
@@ -18,4 +20,7 @@ export SPLUNK_REDIRECT_URL=<REDACTED>
 # CLIENT CREDENTIAL FLOW
 export SPLUNK_APP_CLIENT_CRED_ID=<REDACTED>
 export SPLUNK_APP_CLIENT_CRED_SECRET=<REDACTED>
+export SPLUNK_APP_CLIENT_CRED_ID_2=<REDACTED>
+export SPLUNK_APP_CLIENT_CRED_SECRET_2=<REDACTED>
+export SPLUNK_TENANT_2=<REDACTED>
 export SPLUNK_SCOPE=<REDACTED>
diff --git a/splunk_sdk/auth/auth_manager.py b/splunk_sdk/auth/auth_manager.py
@@ -161,8 +161,11 @@ class AuthManager(ABC):
     the flow that you need for your app.
     """
 
-    def __init__(self, host, client_id, requests_hooks=None):
+    def __init__(self, host, client_id, tenant=None, tenant_scoped=False, region=None, requests_hooks=None):
         self._host = host
+        self._tenant = tenant
+        self._tenant_scoped = tenant_scoped
+        self._region = region
         self._client_id = client_id
         self._context = None
         self._requests_hooks = requests_hooks or []
@@ -190,6 +193,13 @@ def _post(self, url, auth=None, headers=None, data=None, cookies=None):
         return response
 
     def _url(self, path):
+        if self._tenant_scoped is True and self._tenant is not None and self._tenant != "system":
+            self._host = self._tenant + "." + self._host
+            # generate tenant scoped token
+            os.path.join("/", self._tenant, path)
+        # ToDo:Region+system scoped tokens will be added once implementation details are finalized in multi-cell
+        # else:
+        print("https://%s%s" % (self._host, path))
         return "https://%s%s" % (self._host, path)
 
     @staticmethod
@@ -240,8 +250,8 @@ class ClientAuthManager(AuthManager):
     """
 
     # TODO: Host can be an optional value since it has a default
-    def __init__(self, host, client_id, client_secret, scope="", requests_hooks=None):
-        super().__init__(host, client_id, requests_hooks=requests_hooks)
+    def __init__(self, host, client_id, client_secret, scope="", tenant=None, tenant_scoped=False, region=None, requests_hooks=None):
+        super().__init__(host, client_id, tenant, tenant_scoped, region, requests_hooks=requests_hooks)
         self._client_secret = client_secret
         self._scope = scope
 
@@ -269,8 +279,9 @@ class PKCEAuthManager(AuthManager):
     and password, the app through the client_id and redirect_uri. For more details, see identity service documentation.
     """
 
-    def __init__(self, host, client_id, redirect_uri, username, password, scope=DEFAULT_REFRESH_SCOPE, requests_hooks=None):
-        super().__init__(host, client_id, requests_hooks=requests_hooks)
+    def __init__(self, host, client_id, redirect_uri, username, password, scope=DEFAULT_REFRESH_SCOPE, tenant=None,
+                 tenant_scoped=False, region=None, requests_hooks=None):
+        super().__init__(host, client_id, tenant, tenant_scoped, region, requests_hooks=requests_hooks)
         self._redirect_uri = redirect_uri
         self._username = username
         self._password = password
@@ -336,6 +347,8 @@ def _authenticate_user(self, username, password, csrfToken, cookies):
         data = {"username": username, "password": password, "csrfToken": csrfToken}
         response = self._post(self._url(PATH_AUTHN), data=json.dumps(data), cookies=cookies)
         if response.status_code != 200:
+            print("RESPONSE")
+            print(response)
             raise AuthnError("Authentication failed.", response)
         result = response.json()
         status = result.get("status", None)
@@ -432,11 +445,12 @@ def authenticate(self) -> AuthContext:
 
 
 class RefreshTokenAuthManager(AuthManager):
-    def __init__(self, client_id, refresh_token, host, scope="openid", requests_hooks=None):
-        super().__init__(host, client_id, requests_hooks=requests_hooks)
+    def __init__(self, client_id, refresh_token, host, scope="openid", tenant=None, tenant_scoped=False, region=None, requests_hooks=None):
+        super().__init__(host, client_id, tenant, tenant_scoped, region, requests_hooks=requests_hooks)
         self._refresh_token = refresh_token
         self._scope = scope
 
+
     def authenticate(self):
         """Authenticate using the OAuth refresh_token grant type."""
         client_id = self._client_id
@@ -454,7 +468,8 @@ def authenticate(self):
 
 
 class ServicePrincipalAuthManager(AuthManager):
-    def __init__(self, host, principal_name, key, kid, algorithm="ES256", **kwargs):
+    def __init__(self, host, principal_name, key, kid, algorithm="ES256", tenant=None,
+                 tenant_scoped=False, region=None, **kwargs):
         """
         Creates an AuthManager that uses Service Principals to authenticate.
 
@@ -463,7 +478,7 @@ def __init__(self, host, principal_name, key, kid, algorithm="ES256", **kwargs):
         kid is the key_id of `key`
         algorithm is the algorithm that generated `key`
         """
-        super().__init__(host=host, client_id=None, **kwargs)
+        super().__init__(host=host, client_id=None, tenant=tenant, tenant_scoped=tenant_scoped, region=region, **kwargs)
         self._principal_name = principal_name
         self._key = key
         self._kid = kid
diff --git a/splunk_sdk/base_client.py b/splunk_sdk/base_client.py
@@ -238,7 +238,13 @@ def build_url(self, route: str, **kwargs) -> str:
         :param kwargs: TODO DOCS
         :return: The full URL.
         """
-        url = self.context.scheme + "://" + self.context.host
+        if self.context.tenant_scoped is True and self.context.tenant != "system" and route.startswith("/system") is False:
+            url = self.context.scheme + "://" + self.context.tenant + "." + self.context.host
+        elif self.context.tenant_scoped is True and route.startswith("/system/") and self.context.region is not None:
+            url = self.context.scheme + "://region-" + self.context.region + "." + self.context.host
+        else:
+            url = self.context.scheme + "://" + self.context.host
+
         if self.context.port is not None and self.context.port != "":
             url += ":" + self.context.port
 
@@ -247,7 +253,6 @@ def build_url(self, route: str, **kwargs) -> str:
         if route.startswith("/system/") is False and omit_tenant is False:
             url += "/" + self.get_tenant()
         url += route
-
         # set any url path vars
         if len(kwargs) > 0:
             url = url.format(**kwargs)
diff --git a/splunk_sdk/common/context.py b/splunk_sdk/common/context.py
@@ -18,16 +18,23 @@ class Context(object):
     Example:
         client = SplunkCloud(Context(tenant="mytenant"), authManager)
         client.identity.validate()
+
+    Configure tenant_scoped of bool type to True if the hostnames are scoped to a specific tenant/region to support
+    multi-cell environments, currently defaults to False
+    Configure Region as the name of the region that the tenant is contained in
+}
     """
 
     DEFAULT_API_HOST = "api.scp.splunk.com"
     DEFAULT_HOST = DEFAULT_API_HOST
 
     def __init__(self, host=DEFAULT_HOST, api_host=DEFAULT_API_HOST,
-                 port=None, scheme="https", tenant='system', debug=False):
+                 port=None, scheme="https", tenant='system', tenant_scoped=False, region=None, debug=False):
         self.host = host
         self.api_host = api_host
         self.port = port
         self.scheme = scheme
         self.tenant = tenant
         self.debug = debug
+        self.tenant_scoped = tenant_scoped
+        self.region = region
diff --git a/test/auth/test_auth_manager.py b/test/auth/test_auth_manager.py
@@ -10,6 +10,7 @@
 from splunk_sdk.auth import PKCEAuthManager, ClientAuthManager
 from test.fixtures import get_auth_manager as pkce_auth_manager  # NOQA
 from test.fixtures import get_client_auth_manager as client_auth_manager  # NOQA
+from test.fixtures import get_client_auth_manager_scoped as client_auth_manager_scoped
 from test.fixtures import get_service_principal_auth_manager as service_principal_auth_manager  # NOQA
 
 
@@ -66,6 +67,10 @@ def test_client_credentials_authenticate(client_auth_manager):
     auth_context = client_auth_manager.authenticate()
     _assert_client_credentials_auth_context(auth_context)
 
+@pytest.mark.usefixtures('client_auth_manager')  # NOQA
+def test_client_credentials_authenticate_scoped(client_auth_manager_scoped):
+    auth_context = client_auth_manager_scoped.authenticate()
+    _assert_client_credentials_auth_context(auth_context)
 
 @pytest.mark.usefixtures('client_auth_manager')  # NOQA
 def test_token_authenticate(client_auth_manager):
@@ -192,7 +197,7 @@ def _assert_client_credentials_auth_context(auth_context, expires_in=43200):
 
 
 def _assert_sp_credentials_auth_context(auth_context):
-    _assert_client_credentials_auth_context(auth_context, 3600)
+    _assert_client_credentials_auth_context(auth_context)
 
 
 @pytest.mark.usefixtures('service_principal_auth_manager')  # NOQA
diff --git a/test/fixtures.py b/test/fixtures.py
@@ -96,6 +96,22 @@ def get_test_client():
     assert (service_client is not None)
     return service_client
 
+@pytest.fixture(scope='session')
+def get_test_client_scoped_hosts():
+    context = Context(host=os.environ.get('SPLUNK_HOST_2'),
+                      api_host=os.environ.get('SPLUNK_HOST_2'),
+                      tenant=os.environ.get('SPLUNK_TENANT_2'),
+                      tenant_scoped=True,
+                      # Uncomment when services are deployed at region scoped
+                      # region=os.environ.get('SPLUNK_REGION'),
+                      debug=os.environ.get(
+                          'SPLUNK_DEBUG', 'false').lower().strip() == 'true')
+
+    # integration tests use pkce by default
+    service_client = get_client(context, _get_client_manager_scoped())
+    assert (service_client is not None)
+    return service_client
+
 @pytest.fixture(scope='session')
 def get_test_client_ml():
     context = Context(host=os.environ.get('SPLUNK_HOST'),
@@ -138,6 +154,11 @@ def get_client_auth_manager():
     assert (auth_manager is not None)
     return auth_manager
 
+@pytest.fixture(scope='session')
+def get_client_auth_manager_scoped():
+    auth_manager = _get_client_manager_scoped()
+    assert (auth_manager is not None)
+    return auth_manager
 
 @pytest.fixture(scope='session')
 def get_service_principal_auth_manager():
@@ -162,10 +183,19 @@ def _get_client_manager():
                                  'SPLUNK_APP_CLIENT_CRED_SECRET'),
                              scope=os.environ.get('SPLUNK_SCOPE'))
 
+def _get_client_manager_scoped():
+    return ClientAuthManager(host=os.environ.get('SPLUNK_AUTH_HOST_2'),
+                             client_id=os.environ.get(
+                                 'SPLUNK_APP_CLIENT_CRED_ID_2'),
+                             client_secret=os.environ.get(
+                                 'SPLUNK_APP_CLIENT_CRED_SECRET_2'),
+                             scope=os.environ.get('SPLUNK_SCOPE'), tenant=os.environ.get('SPLUNK_TENANT_2'), tenant_scoped=True, region=os.environ.get(
+                                 'SPLUNK_REGION'))
 
 def _get_principal_manager():
     return ServicePrincipalAuthManager(host=os.environ.get('SPLUNK_AUTH_HOST'),
                                        principal_name=os.environ.get('SPLUNK_APP_PRINCIPAL_NAME'),
                                        key=os.environ.get('SPLUNK_APP_PRINCIPAL_PRIVATE_KEY'),
                                        kid=os.environ.get('SPLUNK_APP_PRINCIPAL_KEY_ID'),
                                        algorithm=os.environ.get('SPLUNK_APP_PRINCIPAL_KEY_ALG'))
+
diff --git a/test/test_base_client.py b/test/test_base_client.py
@@ -12,15 +12,15 @@
 import pytest
 
 from test.auth.test_auth_manager import \
-    _assert_client_credentials_auth_context, _assert_pkce_auth_context, _assert_sp_credentials_auth_context  # NOQA
+     _assert_client_credentials_auth_context, _assert_pkce_auth_context, _assert_sp_credentials_auth_context  # NOQA
 from test.fixtures import get_auth_manager as pkce_auth_manager  # NOQA
 from test.fixtures import get_client_auth_manager as client_auth_manager  # NOQA
+from test.fixtures import get_client_auth_manager_scoped as client_auth_manager_scoped  # NOQA
 from test.fixtures import get_service_principal_auth_manager as service_principal_auth_manager  # NOQA
 from splunk_sdk.common.context import Context
 from splunk_sdk.base_client import BaseClient
 from splunk_sdk.identity import Identity as IdentityAndAccessControl
 
-
 @pytest.mark.usefixtures("pkce_auth_manager")  # NOQA
 def test_base_client_instance_with_pkce_auth(pkce_auth_manager):
     """Get a base client with a context and a pkce auth manager."""
@@ -65,7 +65,6 @@ def test_base_client_empty_hooks_with_pkce_auth(pkce_auth_manager):
     IdentityAndAccessControl(base_client).validate_token()
     assert True
 
-
 @pytest.mark.usefixtures("pkce_auth_manager")  # NOQA
 def test_base_client_no_list_hooks_with_pkce_auth(pkce_auth_manager):
     def test_hook(*args, **kwargs):
@@ -85,6 +84,105 @@ def test_base_client_instance_with_client_auth(client_auth_manager):
     assert (base_client is not None)
     _assert_client_credentials_auth_context(base_client.auth_manager.context)
 
+@pytest.mark.usefixtures("client_auth_manager_scoped")  # NOQA
+def test_base_client_instance_with_client_auth_tenant_scoped(client_auth_manager_scoped):
+    """Get a base client with a context and a client auth manager."""
+    context = Context(tenant_scoped=True,
+                      tenant=os.environ.get('SPLUNK_TENANT_2'),
+                      host=os.environ.get('SPLUNK_HOST_2'),
+                      api_host=os.environ.get('SPLUNK_HOST_2'),
+                      region=os.environ.get('SPLUNK_REGION'),
+                      debug=os.environ.get(
+                          'SPLUNK_DEBUG', 'false').lower().strip() == 'true')
+
+    base_client = BaseClient(context=context,
+                             auth_manager=client_auth_manager_scoped)
+    assert (base_client is not None)
+    assert (base_client.context.tenant_scoped is True)
+    assert (base_client.context.tenant == os.environ.get('SPLUNK_TENANT_2'))
+    assert (base_client.context.region == os.environ.get('SPLUNK_REGION'))
+
+    tenant_path = "/exampleservice/path"
+    system_path = "/system/exampleservice/path"
+
+    url_tenant = base_client.build_url(tenant_path)
+    expected_tenant_url_path_template = 'https://%s.%s/%s%s'
+    expected_system_url_path_template = 'https://region-%s.%s%s'
+
+    expected_tenant_url_path = expected_tenant_url_path_template % (os.environ.get('SPLUNK_TENANT_2'), os.environ.get('SPLUNK_HOST_2'), os.environ.get('SPLUNK_TENANT_2'), tenant_path)
+    assert (url_tenant == expected_tenant_url_path)
+
+    expected_system_url_path = expected_system_url_path_template % (os.environ.get('SPLUNK_REGION'), os.environ.get('SPLUNK_HOST_2'), system_path)
+    url_system = base_client.build_url(system_path)
+    assert (url_system == expected_system_url_path)
+
+    _assert_client_credentials_auth_context(base_client.auth_manager.context)
+
+@pytest.mark.usefixtures("client_auth_manager_scoped")  # NOQA
+def test_base_client_instance_with_client_auth_scoped_off(client_auth_manager):
+    """Get a base client with a context and a client auth manager."""
+    context = Context(tenant_scoped=False,
+                      tenant=os.environ.get('SPLUNK_TENANT'),
+                      host=os.environ.get('SPLUNK_HOST'),
+                      api_host=os.environ.get('SPLUNK_HOST'),
+                      region=os.environ.get('SPLUNK_REGION'),
+                      debug=os.environ.get(
+                          'SPLUNK_DEBUG', 'false').lower().strip() == 'true')
+
+    base_client = BaseClient(context=context,
+                             auth_manager=client_auth_manager)
+    assert (base_client is not None)
+    assert (base_client.context.tenant_scoped is False)
+    assert (base_client.context.tenant == os.environ.get('SPLUNK_TENANT'))
+    assert (base_client.context.region == os.environ.get('SPLUNK_REGION'))
+
+    tenant_path = "/exampleservice/path"
+    system_path = "/system/exampleservice/path"
+
+    url_tenant = base_client.build_url(tenant_path)
+    expected_tenant_url_path_template = 'https://%s/%s%s'
+    expected_system_url_path_template = 'https://%s%s'
+
+    expected_tenant_url_path = expected_tenant_url_path_template % (os.environ.get('SPLUNK_HOST'), os.environ.get('SPLUNK_TENANT'), tenant_path)
+    assert (url_tenant == expected_tenant_url_path)
+
+    expected_system_url_path = expected_system_url_path_template % (os.environ.get('SPLUNK_HOST'), system_path)
+    url_system = base_client.build_url(system_path)
+    assert (url_system == expected_system_url_path)
+
+    _assert_client_credentials_auth_context(base_client.auth_manager.context)
+
+@pytest.mark.usefixtures("client_auth_manager_scoped")  # NOQA
+def test_base_client_instance_with_client_auth_default(client_auth_manager):
+    """Get a base client with a context and a client auth manager."""
+    context = Context(tenant=os.environ.get('SPLUNK_TENANT'),
+                      host=os.environ.get('SPLUNK_HOST'),
+                      api_host=os.environ.get('SPLUNK_HOST'),
+                      debug=os.environ.get(
+                          'SPLUNK_DEBUG', 'false').lower().strip() == 'true')
+
+    base_client = BaseClient(context=context,
+                             auth_manager=client_auth_manager)
+    assert (base_client is not None)
+    assert (base_client.context.tenant_scoped is False)
+    assert (base_client.context.tenant == os.environ.get('SPLUNK_TENANT'))
+    assert (base_client.context.region is None)
+
+    tenant_path = "/exampleservice/path"
+    system_path = "/system/exampleservice/path"
+
+    url_tenant = base_client.build_url(tenant_path)
+    expected_tenant_url_path_template = 'https://%s/%s%s'
+    expected_system_url_path_template = 'https://%s%s'
+
+    expected_tenant_url_path = expected_tenant_url_path_template % (os.environ.get('SPLUNK_HOST'), os.environ.get('SPLUNK_TENANT'), tenant_path)
+    assert (url_tenant == expected_tenant_url_path)
+
+    expected_system_url_path = expected_system_url_path_template % (os.environ.get('SPLUNK_HOST'), system_path)
+    url_system = base_client.build_url(system_path)
+    assert (url_system == expected_system_url_path)
+
+    _assert_client_credentials_auth_context(base_client.auth_manager.context)
 
 @pytest.mark.usefixtures("service_principal_auth_manager")  # NOQA
 def test_base_client_instance_with_sp_auth(service_principal_auth_manager):
diff --git a/test/test_forwarders.py b/test/test_forwarders.py
@@ -19,7 +19,7 @@
 
 from splunk_sdk.base_client import BaseClient
 from splunk_sdk.forwarders import SplunkForwarderService, Certificate
-from test.fixtures import get_test_client as test_client  # NOQA
+from test.fixtures import get_test_client_scoped_hosts as test_client  # NOQA
 
 
 @pytest.mark.usefixtures("test_client")  # NOQA
diff --git a/tox.ini b/tox.ini
