diff --git a/drydock_provisioner/cli/commands.py b/drydock_provisioner/cli/commands.py index 400a65e4..fda417b2 100644 --- a/drydock_provisioner/cli/commands.py +++ b/drydock_provisioner/cli/commands.py @@ -103,11 +103,8 @@ def drydock(ctx, debug, url, os_project_domain_name, os_user_domain_name, # setup the drydock client using the passed parameters. url_parse_result = urlparse(url) - if not os_token: - token = KeystoneClient.get_token(ks_sess=ks_sess) - logger.debug("Creating Drydock client with token %s." % token) - else: - token = os_token + def auth_gen(): + return list(ks_sess.get_auth_headers().items()) if not url_parse_result.scheme: ctx.fail('URL must specify a scheme and hostname, optionally a port') @@ -116,7 +113,7 @@ def drydock(ctx, debug, url, os_project_domain_name, os_user_domain_name, scheme=url_parse_result.scheme, host=url_parse_result.hostname, port=url_parse_result.port, - token=token)) + auth_gen=auth_gen)) drydock.add_command(task.task) diff --git a/drydock_provisioner/drydock_client/session.py b/drydock_provisioner/drydock_client/session.py index 7b0c80c2..5ff89a32 100644 --- a/drydock_provisioner/drydock_client/session.py +++ b/drydock_provisioner/drydock_client/session.py @@ -24,17 +24,24 @@ class DrydockSession(object): :param string host: The Drydock server hostname or IP :param int port: (optional) The service port appended if specified - :param string token: Auth token + :param function auth_gen: Callable that will generate a list of authentication + header names and values (2 part tuple) :param string marker: (optional) external context marker """ - def __init__(self, host, port=None, scheme='http', token=None, + def __init__(self, host, port=None, scheme='http', auth_gen=None, marker=None): + self.logger = logging.getLogger(__name__) self.__session = requests.Session() + self.auth_gen = auth_gen + + self.set_auth() + + self.marker = marker self.__session.headers.update({ - 'X-Auth-Token': token, 'X-Context-Marker': marker }) + self.host = host self.scheme = scheme @@ -46,10 +53,14 @@ class DrydockSession(object): # assume default port for scheme self.base_url = "%s://%s/api/" % (self.scheme, self.host) - self.token = token - self.marker = marker - - self.logger = logging.getLogger(__name__) + def set_auth(self): + """Set the session's auth header.""" + if self.auth_gen: + self.logger.debug("Updating session authentication header.") + auth_header = self.auth_gen() + self.__session.headers.update(auth_header) + else: + self.logger.debug("Cannot set auth header, no generator defined.") def get(self, endpoint, query=None): """ @@ -59,8 +70,16 @@ class DrydockSession(object): :param dict query: A dict of k, v pairs to add to the query string :return: A requests.Response object """ - resp = self.__session.get( - self.base_url + endpoint, params=query, timeout=10) + auth_refresh = False + while True: + resp = self.__session.get( + self.base_url + endpoint, params=query, timeout=10) + + if resp.status_code == 401 and not auth_refresh: + self.set_auth() + auth_refresh = True + else: + break return resp @@ -75,16 +94,23 @@ class DrydockSession(object): :param data: Something json.dumps(s) can serialize. Result will be used as the request body :return: A requests.Response object """ + auth_refresh = False + while True: + self.logger.debug("Sending POST with drydock_client session") + if body is not None: + self.logger.debug("Sending POST with explicit body: \n%s" % body) + resp = self.__session.post( + self.base_url + endpoint, params=query, data=body, timeout=10) + else: + self.logger.debug("Sending POST with JSON body: \n%s" % str(data)) + resp = self.__session.post( + self.base_url + endpoint, params=query, json=data, timeout=10) - self.logger.debug("Sending POST with drydock_client session") - if body is not None: - self.logger.debug("Sending POST with explicit body: \n%s" % body) - resp = self.__session.post( - self.base_url + endpoint, params=query, data=body, timeout=10) - else: - self.logger.debug("Sending POST with JSON body: \n%s" % str(data)) - resp = self.__session.post( - self.base_url + endpoint, params=query, json=data, timeout=10) + if resp.status_code == 401 and not auth_refresh: + self.set_auth() + auth_refresh = True + else: + break return resp diff --git a/tests/unit/test_drydock_client.py b/tests/unit/test_drydock_client.py index e764ad99..061f0f93 100644 --- a/tests/unit/test_drydock_client.py +++ b/tests/unit/test_drydock_client.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import pytest +import mock import responses import drydock_provisioner.drydock_client.session as dc_session @@ -40,38 +41,6 @@ def test_session_init_minimal_no_port(): assert dd_ses.base_url == "http://%s/api/" % (host) -def test_session_init_uuid_token(): - host = 'foo.bar.baz' - token = '5f1e08b6-38ec-4a99-9d0f-00d29c4e325b' - - dd_ses = dc_session.DrydockSession(host, token=token) - - assert dd_ses.base_url == "http://%s/api/" % (host) - assert dd_ses.token == token - - -def test_session_init_fernet_token(): - host = 'foo.bar.baz' - token = 'gAAAAABU7roWGiCuOvgFcckec-0ytpGnMZDBLG9hA7Hr9qfvdZDHjsak39YN98HXxoYLIqVm' \ - '19Egku5YR3wyI7heVrOmPNEtmr-fIM1rtahudEdEAPM4HCiMrBmiA1Lw6SU8jc2rPLC7FK7n' \ - 'BCia_BGhG17NVHuQu0S7waA306jyKNhHwUnpsBQ' - - dd_ses = dc_session.DrydockSession(host, token=token) - - assert dd_ses.base_url == "http://%s/api/" % (host) - assert dd_ses.token == token - - -def test_session_init_marker(): - host = 'foo.bar.baz' - marker = '5f1e08b6-38ec-4a99-9d0f-00d29c4e325b' - - dd_ses = dc_session.DrydockSession(host, marker=marker) - - assert dd_ses.base_url == "http://%s/api/" % (host) - assert dd_ses.marker == marker - - @responses.activate def test_session_get(): responses.add( @@ -83,7 +52,10 @@ def test_session_get(): token = '5f1e08b6-38ec-4a99-9d0f-00d29c4e325b' marker = '40c3eaf6-6a8a-11e7-a4bd-080027ef795a' - dd_ses = dc_session.DrydockSession(host, token=token, marker=marker) + def auth_gen(): + return [('X-Auth-Token', token)] + + dd_ses = dc_session.DrydockSession(host, auth_gen=auth_gen, marker=marker) resp = dd_ses.get('v1.0/test') req = resp.request @@ -92,6 +64,31 @@ def test_session_get(): assert req.headers.get('X-Context-Marker', None) == marker +@responses.activate +@mock.patch.object(dc_session.KeystoneClient, 'get_token', + return_value='5f1e08b6-38ec-4a99-9d0f-00d29c4e325b') +def test_session_get_returns_401(*args): + responses.add( + responses.GET, + 'http://foo.bar.baz/api/v1.0/test', + body='okay', + status=401) + host = 'foo.bar.baz' + token = '5f1e08b6-38ec-4a99-9d0f-00d29c4e325b' + marker = '40c3eaf6-6a8a-11e7-a4bd-080027ef795a' + + def auth_gen(): + return [('X-Auth-Token', dc_session.KeystoneClient.get_token())] + + dd_ses = dc_session.DrydockSession(host, auth_gen=auth_gen, marker=marker) + + resp = dd_ses.get('v1.0/test') + req = resp.request + + assert req.headers.get('X-Auth-Token', None) == token + assert req.headers.get('X-Context-Marker', None) == marker + assert dc_session.KeystoneClient.get_token.call_count == 2 + @responses.activate def test_client_task_get(): task = {