diff --git a/pypuppetdb/api.py b/pypuppetdb/api.py index 0d6c70df..529b0c72 100644 --- a/pypuppetdb/api.py +++ b/pypuppetdb/api.py @@ -1,6 +1,7 @@ from __future__ import unicode_literals from __future__ import absolute_import +import hashlib import json import logging import requests @@ -47,6 +48,7 @@ 'pql': 'pdb/query/v4', 'inventory': 'pdb/query/v4/inventory', 'status': 'status/v1/services/puppetdb-status', + 'cmd': 'pdb/cmd/v1' } PARAMETERS = { @@ -58,6 +60,13 @@ 'server_time': 'server_time', } +COMMAND_VERSION = { + "deactivate node": 3, + "replace catalog": 9, + "replace facts": 5, + "store report": 8 +} + ERROR_STRINGS = { 'timeout': 'Connection to PuppetDB timed out on', 'refused': 'Could not reach PuppetDB on', @@ -385,6 +394,79 @@ def _query(self, endpoint, path=None, query=None, self.protocol.upper())) raise + def _cmd(self, command, payload): + """This method posts commands to PuppetDB. Provided a command and payload + it will fire a request at PuppetDB. If PuppetDB can be reached and + answers within the timeout we'll decode the response and give it back + or raise for the HTTP Status Code yesPuppetDB gave back. + + :param command: The PuppetDB Command we want to execute. + :type command: :obj:`string` + :param command: The payload, in wire format, specific to the command. + :type path: :obj:`dict` + + :raises: :class:`~pypuppetdb.errors.EmptyResponseError` + + :returns: The decoded response from PuppetDB + :rtype: :obj:`dict` or :obj:`list` + """ + log.debug('_cmd called with command: {0}, data: {1}'.format( + command, payload)) + + url = self._url('cmd') + + if command not in COMMAND_VERSION: + log.error("Only {0} supported, {1} unsupported".format( + list(COMMAND_VERSION.keys()), command)) + raise APIError + + params = { + "command": command, + "version": COMMAND_VERSION[command], + "certname": payload['certname'], + "checksum": hashlib.sha1(str(payload) # nosec + .encode('utf-8')).hexdigest() + } + + if not self.token: + auth = (self.username, self.password) + else: + auth = None + + try: + r = self._session.post(url, + params=params, + data=json.dumps(payload, default=str), + verify=self.ssl_verify, + cert=(self.ssl_cert, self.ssl_key), + timeout=self.timeout, + auth=auth) + + r.raise_for_status() + + json_body = r.json() + if json_body is not None: + return json_body + else: + del json_body + raise EmptyResponseError + + except requests.exceptions.Timeout: + log.error("{0} {1}:{2} over {3}.".format(ERROR_STRINGS['timeout'], + self.host, self.port, + self.protocol.upper())) + raise + except requests.exceptions.ConnectionError: + log.error("{0} {1}:{2} over {3}.".format(ERROR_STRINGS['refused'], + self.host, self.port, + self.protocol.upper())) + raise + except requests.exceptions.HTTPError as err: + log.error("{0} {1}:{2} over {3}.".format(err.response.text, + self.host, self.port, + self.protocol.upper())) + raise + # Method stubs def nodes(self, unreported=2, with_status=False, **kwargs): @@ -863,3 +945,6 @@ def status(self): :rtype: :obj:`dict` """ return self._query('status') + + def command(self, command, payload): + return self._cmd(command, payload) diff --git a/tests/test_baseapi.py b/tests/test_baseapi.py index cbc9827b..ea2dff5d 100644 --- a/tests/test_baseapi.py +++ b/tests/test_baseapi.py @@ -176,18 +176,22 @@ def test_quote(self, baseapi): + 'facts/macaddress/02%3A42%3Aec%3A94%3A80%3Af0' -class TesteAPIQuery(object): +class TestAPIQuery(object): @mock.patch.object(requests.Session, 'request') def test_timeout(self, get, baseapi): get.side_effect = requests.exceptions.Timeout with pytest.raises(requests.exceptions.Timeout): baseapi._query('nodes') + with pytest.raises(requests.exceptions.Timeout): + baseapi._cmd('deactivate node', {'certname': ''}) @mock.patch.object(requests.Session, 'request') def test_connectionerror(self, get, baseapi): get.side_effect = requests.exceptions.ConnectionError with pytest.raises(requests.exceptions.ConnectionError): baseapi._query('nodes') + with pytest.raises(requests.exceptions.ConnectionError): + baseapi._cmd('deactivate node', {'certname': ''}) @mock.patch.object(requests.Session, 'request') def test_httperror(self, get, baseapi): @@ -195,6 +199,8 @@ def test_httperror(self, get, baseapi): response=requests.Response()) with pytest.raises(requests.exceptions.HTTPError): baseapi._query('nodes') + with pytest.raises(requests.exceptions.HTTPError): + baseapi._cmd('deactivate node', {'certname': ''}) def test_setting_headers_without_token(self, baseapi): httpretty.enable() @@ -386,6 +392,43 @@ def test_query_with_post(self, baseapi, query): httpretty.disable() httpretty.reset() + def test_cmd(self, baseapi, query): + httpretty.reset() + httpretty.enable() + stub_request('http://localhost:8080/pdb/cmd/v1', + method=httpretty.POST) + node_name = 'testnode' + baseapi._cmd('deactivate node', {'certname': node_name}) + last_request = httpretty.last_request() + assert last_request.querystring == { + "certname": [node_name], + "command": ['deactivate node'], + "version": ['3'], + "checksum": ['b93d474970e54943aec050ee399dfb85d21e143a'] + } + assert last_request.headers['Content-Type'] == 'application/json' + assert last_request.method == 'POST' + assert last_request.body == six.b(json.dumps({'certname': node_name})) + httpretty.disable() + httpretty.reset() + + def test_cmd_bad_command(self, baseapi): + httpretty.enable() + stub_request('http://localhost:8080/pdb/cmd/v1') + with pytest.raises(pypuppetdb.errors.APIError): + baseapi._cmd('incorrect command', {}) + httpretty.disable() + httpretty.reset() + + def test_cmd_with_token_authorization(self, token_baseapi): + httpretty.enable() + stub_request('https://localhost:8080/pdb/cmd/v1', + method=httpretty.POST) + token_baseapi._cmd('deactivate node', {'certname': ''}) + assert httpretty.last_request().path.startswith('/pdb/cmd/v1') + assert httpretty.last_request().headers['X-Authentication'] == \ + 'tokenstring' + class TestAPIMethods(object): def test_metric(self, baseapi): @@ -469,3 +512,16 @@ def test_status(self, baseapi): '/status/v1/services/puppetdb-status' httpretty.disable() httpretty.reset() + + def test_command(self, baseapi): + httpretty.enable() + stub_request( + 'http://localhost:8080/pdb/cmd/v1', + method=httpretty.POST + ) + baseapi.command('deactivate node', {'certname': 'testnode'}) + assert httpretty.last_request().path.startswith( + '/pdb/cmd/v1' + ) + httpretty.disable() + httpretty.reset()