Add CSP authentication support (#111)

Signed-off-by: John Cornish <[email protected]>
Co-authored-by: Anil Kodali <[email protected]>
Co-authored-by: Yuqi Jin <[email protected]>
Co-authored-by: John Cornish <[email protected]>

Author: 4 people

.gitignore

@@ -8,3 +8,5 @@ lib/
 lib64/
 env/
 venv/
+.idea/
+.vscode/

example.py

@@ -82,9 +82,26 @@ def send_event(wavefront_client):
 def main():
     """Send sample metrics in a loop."""
     wavefront_proxy_url = sys.argv[1]
+    csp_base_url, csp_api_token = None, None
+    csp_app_id, csp_app_secret, csp_org_id = None, None, None
+
+    if len(sys.argv) == 4:
+        csp_base_url = sys.argv[2]
+        csp_api_token = sys.argv[3]
+    elif len(sys.argv) > 4:
+        csp_base_url = sys.argv[2]
+        csp_app_id = sys.argv[3]
+        csp_app_secret = sys.argv[4]
+        if len(sys.argv) > 5:
+            csp_org_id = sys.argv[5]
 
     client_factory = WavefrontClientFactory()
-    client_factory.add_client(wavefront_proxy_url)
+    client_factory.add_client(wavefront_proxy_url,
+                              csp_base_url=csp_base_url,
+                              csp_api_token=csp_api_token,
+                              csp_app_id=csp_app_id,
+                              csp_app_secret=csp_app_secret,
+                              csp_org_id=csp_org_id)
     wfront_client = client_factory.get_client()
 
     try:
@@ -104,4 +121,16 @@ def main():
     # Either "proxy://our.proxy.lb.com:2878"
     #     Or "https://[email protected]"
     #     should be passed as an input in sys.argv[1]
+    # Or "https://DOMAIN.wavefront.com" "https://CSP_ENDPOINT.cloud.vmware.com"
+    # "CSP_API_TOKEN"
+    #     should be passed as an input in sys.argv[1], sys.argv[2],
+    # and sys.argv[3]
+    # Or "https://DOMAIN.wavefront.com" "https://CSP_ENDPOINT.cloud.vmware.com"
+    # "CSP_APP_ID" "CSP_APP_SECRET"
+    #     should be passed as an input in sys.argv[1], sys.argv[2],
+    # sys.argv[3], and sys.argv[4]
+    # Or "https://DOMAIN.wavefront.com" "https://CSP_ENDPOINT.cloud.vmware.com"
+    # "CSP_APP_ID" "CSP_APP_SECRET" "ORG_ID"
+    #     should be passed as an input in sys.argv[1], sys.argv[2],
+    # sys.argv[3], sys.argv[4], and sys.argv[5]
     main()

setup.py

@@ -11,7 +11,7 @@
 
 setuptools.setup(
     name='wavefront-sdk-python',
-    version='2.0.1',  # Please update with each pull request.
+    version='2.1.0',  # Please update with each pull request.
     author='VMware Aria Operations for Applications Team',
     url='https://github.com/wavefrontHQ/wavefront-sdk-python',
     license='Apache-2.0',

test/test_client.py

@@ -5,6 +5,7 @@
 
 import unittest
 import uuid
+from unittest.mock import Mock
 
 import requests
 
@@ -23,7 +24,7 @@ def setUp(self):
             enable_internal_metrics=False)
         self._spans_log_buffer = self._sender._spans_log_buffer
         self._tracing_spans_buffer = self._sender._tracing_spans_buffer
-        self._response = unittest.mock.Mock()
+        self._response = Mock()
         self._response.status_code = 200
 
     def test_send_version_with_internal_metrics(self):
@@ -119,7 +120,7 @@ def test_report_event(self):
             self._sender._report('',
                                  self._sender.WAVEFRONT_EVENT_FORMAT,
                                  '',
-                                 unittest.mock.Mock())
+                                 Mock())
             mock_post.assert_called_once_with(unittest.mock.ANY,
                                               headers=unittest.mock.ANY,
                                               data=unittest.mock.ANY)
@@ -129,7 +130,7 @@ def test_report_non_event(self):
         with unittest.mock.patch.object(
                 requests.Session, 'post',
                 return_value=self._response) as mock_post:
-            self._sender._report('', 'metric', '', unittest.mock.Mock())
+            self._sender._report('', 'metric', '', Mock())
             mock_post.assert_called_once_with(unittest.mock.ANY,
                                               params=unittest.mock.ANY,
                                               headers=unittest.mock.ANY,

test/test_wavefront_client.py

@@ -57,6 +57,41 @@ def test_get_client(self):
                                    wavefront_sdk.multi_clients
                                    .WavefrontMultiClient))
 
+    def test_get_csp_client(self):
+        """Test get_client of WavefrontClientFactory
+        for client with CSP configuration
+        """
+
+        wavefront_cluster_base_url = "https://csp-enabled.wavefront.com"
+
+        # if there is only one client and using CSP OAuth App ID and App Secret
+        fake_oauth_app_id = "fake-csp-app-id"
+        fake_oauth_app_secret = "fake-csp-app-secret"
+        multi_client_factory = WavefrontClientFactory()
+        multi_client_factory.add_client(
+            wavefront_cluster_base_url,
+            csp_app_id=fake_oauth_app_id,
+            csp_app_secret=fake_oauth_app_secret,
+        )
+        self.assertTrue(isinstance(multi_client_factory.get_client(),
+                                   wavefront_sdk.client.WavefrontClient))
+        self.assertTrue(
+            multi_client_factory.existing_client(wavefront_cluster_base_url)
+        )
+
+        # if there is only one client and using CSP Api Token
+        fake_csp_api_token = "fake-csp-api-token"
+        multi_client_factory = WavefrontClientFactory()
+        multi_client_factory.add_client(
+            wavefront_cluster_base_url,
+            csp_api_token=fake_csp_api_token,
+        )
+        self.assertTrue(isinstance(multi_client_factory.get_client(),
+                                   wavefront_sdk.client.WavefrontClient))
+        self.assertTrue(
+            multi_client_factory.existing_client(wavefront_cluster_base_url)
+        )
+
     def test_existing_client(self):
         """Test existing_client of WavefrontClientFactory"""
         server = "http://192.114.71.230:2878"

wavefront_sdk/auth/__init__.py

@@ -0,0 +1,16 @@
+"""Wavefront SDK Authentication.
+
+@author Jerry Belmonte ([email protected])
+"""
+
+from .csp.authorize_response import AuthorizeResponse
+from .csp.csp_request import CSPAPIToken, CSPClientCredentials
+from .csp.csp_token_service import CSPAccessTokenService
+from .csp.token_service_factory import TokenServiceProvider
+
+
+__all__ = ['AuthorizeResponse',
+           'CSPAPIToken',
+           'CSPClientCredentials',
+           'CSPAccessTokenService',
+           'TokenServiceProvider']

wavefront_sdk/auth/csp/__init__.py

@@ -0,0 +1,4 @@
+"""CSP Authentication.
+
+@author Jerry Belmonte ([email protected])
+"""

wavefront_sdk/auth/csp/authorize_response.py

@@ -0,0 +1,37 @@
+"""CSP Authorize Response.
+
+@author Jerry Belmonte ([email protected])
+"""
+
+from dataclasses import dataclass, field
+
+
+@dataclass
+class AuthorizeResponse:
+    """Authorize Response."""
+
+    access_token: str = field(default='', repr=False)
+    expires_in: int = 0
+
+    def set_auth_response(self, response):
+        """Set the CSP auth response.
+
+        @param response: The json-encoded response.
+        """
+        self.access_token = response.get("access_token")
+        self.expires_in = response.get("expires_in")
+
+    def get_time_offset(self):
+        """Calculate the time offset.
+
+        Calculates the time offset for scheduling regular requests to CSP based
+        on the expiration time of the token. If the access token expiration
+        time is less than 10 minutes, schedule requests 30 seconds before it
+        expires. If the access token expiration time is 10 minutes or more,
+        schedule requests 3 minutes before it expires.
+
+        @return: The expiration time offset of the CSP access token in seconds.
+        """
+        if self.expires_in < 600:
+            return self.expires_in - 30
+        return self.expires_in - 180

wavefront_sdk/auth/csp/csp_request.py

@@ -0,0 +1,80 @@
+"""CSP Request types.
+
+@author Jerry Belmonte ([email protected])
+"""
+
+from base64 import b64encode
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class AuthServerURL:
+    """Auth Server URL."""
+
+    base_url: str
+    auth_path: str
+
+    def get_server_url(self):
+        """Get the full authentication server URL.
+
+        @return: The authentication server URL.
+        """
+        if self.base_url.endswith('/'):
+            return self.base_url[:-1] + self.auth_path
+        return self.base_url + self.auth_path
+
+
+@dataclass(frozen=True)
+class CSPAPIToken(AuthServerURL):
+    """CSP Api Token."""
+
+    token: str
+
+    def get_data(self):
+        """Get the HTTP request body.
+
+        @return: The data for the request body.
+        """
+        return {'api_token': self.token}
+
+    def get_headers(self):
+        """Get the HTTP request headers.
+
+        @return: The parameters for request headers.
+        """
+        return {'Content-Type': 'application/x-www-form-urlencoded'}
+
+
+@dataclass(frozen=True)
+class CSPClientCredentials(AuthServerURL):
+    """CSP Client Credentials."""
+
+    client_id: str
+    client_secret: str
+    org_id: str = ''
+
+    def get_data(self):
+        """Get the HTTP request body.
+
+        @return: The data for the request body.
+        """
+        data = {'grant_type': 'client_credentials'}
+        if self.org_id:
+            data['orgId'] = self.org_id
+        return data
+
+    def encode_csp_credentials(self):
+        """Encode the CSP client credentials.
+
+        @return: Base64 encoded client credentials.
+        """
+        csp_credentials = self.client_id + ":" + self.client_secret
+        return b64encode(csp_credentials.encode("utf-8")).decode("utf-8")
+
+    def get_headers(self):
+        """Get the HTTP request headers.
+
+        @return: The parameters for request headers.
+        """
+        return {'Authorization': f'Basic {self.encode_csp_credentials()}',
+                'Content-Type': 'application/x-www-form-urlencoded'}

wavefront_sdk/auth/csp/csp_token_service.py

@@ -0,0 +1,86 @@
+"""CSP Token Service implementation.
+
+@author Jerry Belmonte ([email protected])
+"""
+
+import logging
+from time import time
+
+from requests import HTTPError, Timeout, post
+
+from .authorize_response import AuthorizeResponse
+from .token_service import TokenService
+
+
+LOGGER = logging.getLogger('wavefront_sdk.auth.csp.CSPAccessTokenService')
+CSP_API_TOKEN_SERVICE_TYPE = 'API_TOKEN'
+CSP_OAUTH_TOKEN_SERVICE_TYPE = 'OAUTH'
+CSP_REQUEST_TIMEOUT_SEC = 30
+
+
+class CSPAccessTokenService(TokenService):
+    """CSP Access Token Service Implementation."""
+
+    def __init__(self, csp_type: str, csp_service):
+        """Construct CSP Access Token Service.
+
+        @param csp_type: The csp token service type.
+        @param csp_service: The csp service object.
+        """
+        self._csp_type = csp_type
+        self._csp_service = csp_service
+        self._csp_access_token = None
+        self._token_expiration_time = 0
+        self._csp_response = None
+
+    def get_type(self):
+        """Get the token service type.
+
+        @return: The service type.
+        """
+        return self._csp_type
+
+    def get_token(self):
+        """Get the token service access token.
+
+        @return: The access token.
+        """
+        LOGGER.debug("Retrieving the CSP access token.")
+        try:
+            response = post(self._csp_service.get_server_url(),
+                            data=self._csp_service.get_data(),
+                            headers=self._csp_service.get_headers(),
+                            timeout=CSP_REQUEST_TIMEOUT_SEC)
+            code = response.status_code
+            if code == 200:
+                self._csp_response = AuthorizeResponse()
+                self._csp_response.set_auth_response(response.json())
+                self._csp_access_token = self._csp_response.access_token
+                self._token_expiration_time =\
+                    time() + self._csp_response.get_time_offset()
+                LOGGER.info("CSP auth token refresh succeeded, access token"
+                            " expires in %d seconds.",
+                            self._csp_response.expires_in)
+                return self._csp_access_token
+            if not response.ok:
+                data = response.json()
+                LOGGER.error("CSP auth token refresh failed with status code:"
+                             " %d %s", code, data.get("message"))
+                response.raise_for_status()
+
+        except HTTPError as error:
+            LOGGER.error("CSP HTTP Error: %s", error)
+        except ConnectionError as error:
+            LOGGER.error("CSP Connection Error: %s", error)
+        except Timeout as error:
+            LOGGER.error("CSP Timeout Error: %s", error)
+        return None
+
+    def get_csp_access_token(self):
+        """Get the CSP access token.
+
+        @return: The access token.
+        """
+        if not self._csp_access_token or time() >= self._token_expiration_time:
+            return self.get_token()
+        return self._csp_access_token

wavefront_sdk/auth/csp/token_service.py

@@ -0,0 +1,24 @@
+"""CSP Token Service.
+
+@author Jerry Belmonte ([email protected])
+"""
+
+from abc import ABC, abstractmethod
+
+
+class TokenService(ABC):
+    """Service that gets access tokens."""
+
+    @abstractmethod
+    def get_token(self) -> str:
+        """Get the token service access token.
+
+        @return: The access token.
+        """
+
+    @abstractmethod
+    def get_type(self) -> str:
+        """Get the token service type.
+
+        @return: The service type.
+        """