Add ability to handle non-JSON endpoints

src/oktacli/cli.py

@@ -2033,6 +2033,13 @@ def get_users_for(obj_list, rest_path, workers=1):
     type=click.Choice(("get", "post", "put", "delete")),
     help="Which HTTP method to use; default: 'get'",
 )
+@click.option(
+    "-H",
+    "--header",
+    "headers",
+    multiple=True,
+    help="Set an HTTP header",
+)
 @click.option(
     "-q",
     "--query",
@@ -2054,7 +2061,7 @@ def get_users_for(obj_list, rest_path, workers=1):
     help="Specify a different base path than the default (/api/v1)",
 )
 @_output_type_command_wrapper(None)
-def raw(api_endpoint, http_method, query_params, body, base_path, **kwargs):
+def raw(api_endpoint, http_method, query_params, body, base_path, headers, **kwargs):
     """Perform a request against the specified API endpoint"""
     methods = {
         "get": REST.get,
@@ -2068,6 +2075,7 @@ def raw(api_endpoint, http_method, query_params, body, base_path, **kwargs):
     if not api_endpoint.startswith("/"):
         api_endpoint = "/" + api_endpoint
     p_dict = dict([(y[0], y[1]) for y in map(lambda x: x.split("=", 1), query_params)])
+    header_dict = dict([(h[0].strip(), h[1].strip()) for h in map(lambda x: x.split(":", 1), headers)])
     if body:
         if body.startswith("FILE:"):
             use_body = json.loads(open(body[5:], "r").read())
@@ -2079,6 +2087,7 @@ def raw(api_endpoint, http_method, query_params, body, base_path, **kwargs):
         api_endpoint,
         use_method,
         params=p_dict,
+        headers=header_dict,
         body_obj=use_body,
         custom_path_base=base_path,
     )

src/oktacli/okta.py

@@ -61,12 +61,16 @@ def call_okta_raw(
         method,
         *,
         params=None,
+        headers=None,
         body_obj=None,
         custom_url=None,
         custom_path_base=None,
     ):
         call_method = getattr(self.session, method.value)
-        call_params = {"params": params if params is not None else {}}
+        call_params = {
+            "params": params if params is not None else {},
+            "headers": headers if headers is not None else {},
+        }
         call_url = urljoin(
             (custom_url if custom_url is not None else self.url),
             "/".join(
@@ -112,6 +116,7 @@ def call_okta(
         method,
         *,
         params=None,
+        headers=None,
         body_obj=None,
         result_limit=None,
         custom_url=None,
@@ -121,37 +126,41 @@ def call_okta(
             path,
             method,
             params=params,
+            headers=headers,
             body_obj=body_obj,
             custom_url=custom_url,
             custom_path_base=custom_path_base,
         )
-        rv = rsp.json()
-        # NOW, we either have a SINGLE DICT in the rv variable,
-        #     *OR*
-        # a list.
-        last_url = None
-        while True:
-            # let's stop if we defined a result_limit
-            if result_limit and len(rv) > result_limit:
-                break
-            # now, let's get all the "next" links. if we do NOT have a list,
-            # we do not have "next" links :) . handy!
-            url = rsp.links.get("next", {"url": ""})["url"]
-            # sanity checks
-            if not url or last_url == url:
-                break
-            last_url = url
-            rsp = self.call_okta_raw("", REST.get, custom_url=url, custom_path_base="")
-            # now the += operation is safe, cause we have a list.
-            # this is a liiiitle bit implicit, but should work smoothly.
-            rv += rsp.json()
-        # filter out _links items from the final result list
-        if isinstance(rv, list):
-            for item in rv:
-                item.pop("_links", None)
-        elif isinstance(rv, dict):
-            rv.pop("_links", None)
-        return rv
+        if rsp.headers['Content-Type'] == "application/json":
+            rv = rsp.json()
+            # NOW, we either have a SINGLE DICT in the rv variable,
+            #     *OR*
+            # a list.
+            last_url = None
+            while True:
+                # let's stop if we defined a result_limit
+                if result_limit and len(rv) > result_limit:
+                    break
+                # now, let's get all the "next" links. if we do NOT have a list,
+                # we do not have "next" links :) . handy!
+                url = rsp.links.get("next", {"url": ""})["url"]
+                # sanity checks
+                if not url or last_url == url:
+                    break
+                last_url = url
+                rsp = self.call_okta_raw("", REST.get, custom_url=url, custom_path_base="")
+                # now the += operation is safe, cause we have a list.
+                # this is a liiiitle bit implicit, but should work smoothly.
+                rv += rsp.json()
+            # filter out _links items from the final result list
+            if isinstance(rv, list):
+                for item in rv:
+                    item.pop("_links", None)
+            elif isinstance(rv, dict):
+                rv.pop("_links", None)
+            return rv
+        else:
+            return rsp.text
 
     def list_groups(self, query_ex="", filter_ex=""):
         params = {}

tests/test_cli_raw.py

@@ -0,0 +1,32 @@
+import re
+from unittest.mock import patch
+
+from click.testing import CliRunner
+import responses
+
+from oktacli import cli
+from oktacli.okta import Okta
+from .testprep import prepare_standard_calls
+
+
+@patch("oktacli.cli.get_manager")
+@responses.activate
+def test_raw_accept_header(get_manager):
+    # test data
+    params0 = ["raw", "-H", "Accept: application/xml", "-H", "Some: header", "some/endpoint"]
+    wanted_result = {"test_add_user_to_group": "ok"}
+    # set up test
+    get_manager.return_value = Okta("http://okta", "12ab")
+    runner = CliRunner()
+    responses.add(
+        responses.GET,
+        re.compile(".+/some/endpoint/?"),
+        body="",
+        status=204,
+        match=[responses.matchers.header_matcher({"Accept": "application/xml", "Some": "header"})],
+    )
+    # run command
+    result = runner.invoke(cli.cli_main, params0)
+    # validate
+    assert result.exit_code == 0, result.stdout
+    assert result.exception is None