Add support for the Command API, /pdb/cmd/v1.

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)

tests/test_baseapi.py

@@ -176,25 +176,31 @@ 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):
         get.side_effect = requests.exceptions.HTTPError(
             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()