From 60f335729f7c7f0dfd26b83c54a31710463b8683 Mon Sep 17 00:00:00 2001 From: Ben Jansen Date: Fri, 5 Jan 2024 09:50:13 -0800 Subject: [PATCH] Add ability to handle non-JSON endpoints Add an option to pass arbitrary HTTP request headers in the `raw` command, for endpoints such as `[app]/sso/saml/metadata` that are not JSON. Also do not attempt to parse the body as JSON if the content type indicates something other than JSON. --- src/oktacli/cli.py | 11 +++++++- src/oktacli/okta.py | 65 ++++++++++++++++++++++++------------------- tests/test_cli_raw.py | 32 +++++++++++++++++++++ 3 files changed, 79 insertions(+), 29 deletions(-) create mode 100644 tests/test_cli_raw.py diff --git a/src/oktacli/cli.py b/src/oktacli/cli.py index d591c19..5726e33 100644 --- a/src/oktacli/cli.py +++ b/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, ) diff --git a/src/oktacli/okta.py b/src/oktacli/okta.py index 2c9a33b..e993fcf 100644 --- a/src/oktacli/okta.py +++ b/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 = {} diff --git a/tests/test_cli_raw.py b/tests/test_cli_raw.py new file mode 100644 index 0000000..20fb9af --- /dev/null +++ b/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