Skip to content

Commit

Permalink
Add OAuth functionality to the HTTP util
Browse files Browse the repository at this point in the history
  • Loading branch information
@ofek
ofek committed Sep 9, 2022
1 parent ac2626a commit ad2f5e7
Show file tree
Hide file tree
Showing 4 changed files with 207 additions and 0 deletions.
3 changes: 3 additions & 0 deletions datadog_checks_base/datadog_checks/base/data/agent_requirements.in
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,8 @@ lz4==2.2.1; python_version < '3.0'
lz4==3.1.3; python_version > '3.0'
mmh3==2.5.1; python_version < '3.0'
mmh3==3.0.0; python_version > '3.0'
oauthlib==3.1.0; python_version < '3.0'
oauthlib==3.2.0; python_version > '3.0'
openstacksdk==0.39.0; python_version < '3.0'
openstacksdk==0.61.0; python_version > '3.0'
orjson==3.7.11; python_version > '3.0'
Expand Down Expand Up @@ -87,6 +89,7 @@ redis==4.3.4; python_version > '3.0'
requests-kerberos==0.12.0; python_version < '3.0'
requests-kerberos==0.14.0; python_version > '3.0'
requests-ntlm==1.1.0
requests-oauthlib==1.3.1
requests-toolbelt==0.9.1
requests-unixsocket==0.3.0
requests==2.27.1; python_version < '3.0'
Expand Down
62 changes: 62 additions & 0 deletions datadog_checks_base/datadog_checks/base/utils/http.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,8 @@
requests_aws = None
requests_kerberos = None
requests_ntlm = None
requests_oauthlib = None
oauth2 = None
jwt = None
default_backend = None
serialization = None
Expand Down Expand Up @@ -794,6 +796,65 @@ def read(self, **request):
return self._token


class AuthTokenOAuthReader(object):
def __init__(self, config):
self._url = config.get('url', '')
if not isinstance(self._url, str):
raise ConfigurationError('The `url` setting of `auth_token` reader must be a string')
elif not self._url:
raise ConfigurationError('The `url` setting of `auth_token` reader is required')

self._client_id = config.get('client_id', '')
if not isinstance(self._client_id, str):
raise ConfigurationError('The `client_id` setting of `auth_token` reader must be a string')
elif not self._client_id:
raise ConfigurationError('The `client_id` setting of `auth_token` reader is required')

self._client_secret = config.get('client_secret', '')
if not isinstance(self._client_secret, str):
raise ConfigurationError('The `client_secret` setting of `auth_token` reader must be a string')
elif not self._client_secret:
raise ConfigurationError('The `client_secret` setting of `auth_token` reader is required')

self._basic_auth = config.get('basic_auth', False)
if not isinstance(self._basic_auth, bool):
raise ConfigurationError('The `basic_auth` setting of `auth_token` reader must be a boolean')

self._fetch_options = {'token_url': self._url}
if self._basic_auth:
self._fetch_options['auth'] = requests_auth.HTTPBasicAuth(self._client_id, self._client_secret)
else:
self._fetch_options['client_id'] = self._client_id
self._fetch_options['client_secret'] = self._client_secret

self._token = None
self._expiration = None

def read(self, **request):
if self._token is None or get_timestamp() >= self._expiration or 'error' in request:
global oauth2
if oauth2 is None:
from oauthlib import oauth2

global requests_oauthlib
if requests_oauthlib is None:
import requests_oauthlib

client = oauth2.BackendApplicationClient(client_id=self._client_id)
oauth = requests_oauthlib.OAuth2Session(client=client)
response = oauth.fetch_token(**self._fetch_options)

# https://www.rfc-editor.org/rfc/rfc6749#section-5.2
if 'error' in response:
raise Exception('OAuth2 client credentials grant error: {}'.format(response['error']))

# https://www.rfc-editor.org/rfc/rfc6749#section-4.4.3
self._token = response['access_token']
self._expiration = get_timestamp() + response['expires_in']

return self._token


class DCOSAuthTokenReader(object):
def __init__(self, config):
self._login_url = config.get('login_url', '')
Expand Down Expand Up @@ -893,6 +954,7 @@ def write(self, token, **request):

AUTH_TOKEN_READERS = {
'file': AuthTokenFileReader,
'oauth': AuthTokenOAuthReader,
'dcos_auth': DCOSAuthTokenReader,
}
AUTH_TOKEN_WRITERS = {'header': AuthTokenHeaderWriter}
Expand Down
3 changes: 3 additions & 0 deletions datadog_checks_base/pyproject.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -75,12 +75,15 @@ http = [
"aws-requests-auth==0.4.3",
"botocore==1.20.112; python_version < '3.0'",
"botocore==1.27.44; python_version > '3.0'",
"oauthlib==3.1.0; python_version < '3.0'",
"oauthlib==3.2.0; python_version > '3.0'",
"pyjwt==1.7.1; python_version < '3.0'",
"pyjwt==2.4.0; python_version > '3.0'",
"pysocks==1.7.1",
"requests-kerberos==0.12.0; python_version < '3.0'",
"requests-kerberos==0.14.0; python_version > '3.0'",
"requests-ntlm==1.1.0",
"requests-oauthlib==1.3.1",
"win-inet-pton==1.1.0; sys_platform == 'win32' and python_version < '3.0'",
]
json = [
Expand Down
139 changes: 139 additions & 0 deletions datadog_checks_base/tests/base/utils/http/test_authtoken.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -141,6 +141,79 @@ def test_pattern_no_groups(self):
RequestsWrapper(instance, init_config)


class TestAuthTokenOAuthReaderCreation:
def test_url_missing(self):
instance = {'auth_token': {'reader': {'type': 'oauth'}, 'writer': {'type': 'header'}}}
init_config = {}

with pytest.raises(ConfigurationError, match='^The `url` setting of `auth_token` reader is required$'):
RequestsWrapper(instance, init_config)

def test_url_not_string(self):
instance = {'auth_token': {'reader': {'type': 'oauth', 'url': {}}, 'writer': {'type': 'header'}}}
init_config = {}

with pytest.raises(ConfigurationError, match='^The `url` setting of `auth_token` reader must be a string$'):
RequestsWrapper(instance, init_config)

def test_client_id_missing(self):
instance = {'auth_token': {'reader': {'type': 'oauth', 'url': 'foo'}, 'writer': {'type': 'header'}}}
init_config = {}

with pytest.raises(ConfigurationError, match='^The `client_id` setting of `auth_token` reader is required$'):
RequestsWrapper(instance, init_config)

def test_client_id_not_string(self):
instance = {
'auth_token': {'reader': {'type': 'oauth', 'url': 'foo', 'client_id': {}}, 'writer': {'type': 'header'}}
}
init_config = {}

with pytest.raises(
ConfigurationError, match='^The `client_id` setting of `auth_token` reader must be a string$'
):
RequestsWrapper(instance, init_config)

def test_client_secret_missing(self):
instance = {
'auth_token': {'reader': {'type': 'oauth', 'url': 'foo', 'client_id': 'bar'}, 'writer': {'type': 'header'}}
}
init_config = {}

with pytest.raises(
ConfigurationError, match='^The `client_secret` setting of `auth_token` reader is required$'
):
RequestsWrapper(instance, init_config)

def test_client_secret_not_string(self):
instance = {
'auth_token': {
'reader': {'type': 'oauth', 'url': 'foo', 'client_id': 'bar', 'client_secret': {}},
'writer': {'type': 'header'},
}
}
init_config = {}

with pytest.raises(
ConfigurationError, match='^The `client_secret` setting of `auth_token` reader must be a string$'
):
RequestsWrapper(instance, init_config)

def test_basic_auth_not_boolean(self):
instance = {
'auth_token': {
'reader': {'type': 'oauth', 'url': 'foo', 'client_id': 'bar', 'client_secret': 'baz', 'basic_auth': {}},
'writer': {'type': 'header'},
}
}
init_config = {}

with pytest.raises(
ConfigurationError, match='^The `basic_auth` setting of `auth_token` reader must be a boolean$'
):
RequestsWrapper(instance, init_config)


class TestAuthTokenDCOSReaderCreation:
def test_login_url_missing(self):
instance = {'auth_token': {'reader': {'type': 'dcos_auth'}, 'writer': {'type': 'header'}}}
Expand Down Expand Up @@ -369,6 +442,72 @@ def test_pattern_match(self):
assert http.options['headers'] == expected_headers


class TestAuthTokenOAuth:
def test_failure(self):
instance = {
'auth_token': {
'reader': {'type': 'oauth', 'url': 'foo', 'client_id': 'bar', 'client_secret': 'baz'},
'writer': {'type': 'header', 'name': 'Authorization', 'value': 'Bearer <TOKEN>'},
}
}
init_config = {}
http = RequestsWrapper(instance, init_config)

expected_headers = {'Authorization': 'Bearer foo'}
expected_headers.update(DEFAULT_OPTIONS['headers'])

class MockOAuth2Session(object):
def __init__(self, *args, **kwargs):
pass

def fetch_token(self, *args, **kwargs):
return {'error': 'unauthorized_client'}

with mock.patch('requests.get'), mock.patch('oauthlib.oauth2.BackendApplicationClient'), mock.patch(
'requests_oauthlib.OAuth2Session', side_effect=MockOAuth2Session
):
with pytest.raises(Exception, match='OAuth2 client credentials grant error: unauthorized_client'):
http.get('https://www.google.com')

def test_success(self):
instance = {
'auth_token': {
'reader': {'type': 'oauth', 'url': 'foo', 'client_id': 'bar', 'client_secret': 'baz'},
'writer': {'type': 'header', 'name': 'Authorization', 'value': 'Bearer <TOKEN>'},
}
}
init_config = {}
http = RequestsWrapper(instance, init_config)

expected_headers = {'Authorization': 'Bearer foo'}
expected_headers.update(DEFAULT_OPTIONS['headers'])

class MockOAuth2Session(object):
def __init__(self, *args, **kwargs):
pass

def fetch_token(self, *args, **kwargs):
return {'access_token': 'foo', 'expires_in': 9000}

with mock.patch('requests.get') as get, mock.patch('oauthlib.oauth2.BackendApplicationClient'), mock.patch(
'requests_oauthlib.OAuth2Session', side_effect=MockOAuth2Session
):
http.get('https://www.google.com')

get.assert_called_with(
'https://www.google.com',
headers=expected_headers,
auth=None,
cert=None,
proxies=None,
timeout=(10.0, 10.0),
verify=True,
allow_redirects=True,
)

assert http.options['headers'] == expected_headers


class TestAuthTokenDCOS:
def test_token_auth(self):
priv_key_path = os.path.join(FIXTURE_PATH, 'dcos', 'private-key.pem')
Expand Down

0 comments on commit ad2f5e7

Please sign in to comment.