From a10e08d91ecf6202ec46c4273bb516099874e225 Mon Sep 17 00:00:00 2001
From: Brad Klein <bklein@cray.com>
Date: Thu, 6 Jul 2023 14:13:40 -0600
Subject: [PATCH] Add tenant name header allowing other headers (CASMINST-6514)

---
 cray/constants.py |  5 +++++
 cray/generator.py |  7 ++++---
 cray/rest.py      | 10 +++++++---
 cray/utils.py     | 11 +----------
 4 files changed, 17 insertions(+), 16 deletions(-)

diff --git a/cray/constants.py b/cray/constants.py
index 423158e..7d7e6fd 100644
--- a/cray/constants.py
+++ b/cray/constants.py
@@ -44,6 +44,8 @@ def _make_envvar(var):
 DANGER_TAG = "cli_danger"
 FROM_FILE_TAG = "cli_from_file"
 CONVERSION_FLAG = "cray_converted"
+HEADER_ORIGIN = "header"
+HEADERS_ORIGIN = "headers"
 
 # Config constants
 DEFAULT_CONFIG = 'default'
@@ -52,3 +54,6 @@ def _make_envvar(var):
 CONFIG_DIR_NAME = 'configurations'
 LOG_DIR_NAME = 'logs'
 AUTH_DIR_NAME = 'tokens'
+
+# Rest constants
+TENANT_HEADER_NAME_KEY = "Cray-Tenant-Name"
diff --git a/cray/generator.py b/cray/generator.py
index beb3d8c..ccad5e7 100644
--- a/cray/generator.py
+++ b/cray/generator.py
@@ -37,6 +37,8 @@
 from cray.constants import CONVERSION_FLAG
 from cray.constants import DANGER_TAG
 from cray.constants import FROM_FILE_TAG
+from cray.constants import HEADER_ORIGIN
+from cray.constants import HEADERS_ORIGIN
 from cray.constants import HIDDEN_TAG
 from cray.constants import IGNORE_TAG
 from cray.constants import TAG_SPLIT
@@ -44,7 +46,6 @@
 
 PATH_ORIGIN = 'path'
 QUERY_ORIGIN = 'query'
-HEADER_ORIGIN = 'header'
 PARAM_ORIGIN = 'params'
 FILE_ORIGIN = 'filepath'
 
@@ -169,7 +170,7 @@ def _parse_data(data, base=None, **kwargs):
     if opts[QUERY_ORIGIN]:
         args['params'] = opts[QUERY_ORIGIN]
     if opts[HEADER_ORIGIN]:
-        args['headers'] = opts[HEADER_ORIGIN]
+        args[HEADERS_ORIGIN] = opts[HEADER_ORIGIN]
     if opts[FILE_ORIGIN]:
         fields = {}
         for k, v in opts[FILE_ORIGIN].items():
@@ -177,7 +178,7 @@ def _parse_data(data, base=None, **kwargs):
             # pylint: disable=consider-using-with
             fields[k] = (os.path.basename(v), open(v, 'rb'))
         args['data'] = MultipartEncoder(fields=fields)
-        args.setdefault('headers', {})['Content-Type'] = args[
+        args.setdefault(HEADERS_ORIGIN, {})['Content-Type'] = args[
             'data'].content_type
     return (method, route, args)
 
diff --git a/cray/rest.py b/cray/rest.py
index fa15b07..00b0743 100644
--- a/cray/rest.py
+++ b/cray/rest.py
@@ -34,6 +34,8 @@
 from six.moves import urllib
 from urllib3.exceptions import InsecureRequestWarning
 
+from cray.constants import HEADERS_ORIGIN
+from cray.constants import TENANT_HEADER_NAME_KEY
 from cray.echo import echo
 from cray.echo import LOG_DEBUG
 from cray.echo import LOG_RAW
@@ -41,7 +43,7 @@
 from cray.errors import InsecureError
 from cray.errors import UnauthorizedError
 from cray.utils import get_hostname
-from cray.utils import get_headers
+from cray.utils import get_tenant
 
 
 def make_url(route, url=None, default_scheme='https', ctx=None):
@@ -91,13 +93,15 @@ def request(method, route, callback=None, **kwargs):
 
     try:
         url = make_url(route)
-        headers = get_headers(ctx=ctx)
+        tenant = get_tenant(ctx=ctx)
+        if tenant:
+            opts.setdefault(HEADERS_ORIGIN, {})[TENANT_HEADER_NAME_KEY] = tenant
         echo(f'REQUEST: {method} to {url}', ctx=ctx, level=LOG_DEBUG)
         echo(f'OPTIONS: {opts}', ctx=ctx, level=LOG_RAW)
         # TODO: Find solution for this.
         with warnings.catch_warnings():
             warnings.filterwarnings("ignore", category=InsecureRequestWarning)
-            response = requester.request(method, url, **opts, headers=headers)
+            response = requester.request(method, url, **opts)
             if not response.ok:
                 _log_request_error(response.text, ctx)
                 raise BadResponseError(response, ctx=ctx)
diff --git a/cray/utils.py b/cray/utils.py
index d3b835c..493a50d 100644
--- a/cray/utils.py
+++ b/cray/utils.py
@@ -78,18 +78,9 @@ def get_tenant(ctx=None):
     """ Get the tenant requests are scoped to (None if not scoped to a tenant) """
     ctx = ctx or click.get_current_context()
     config = ctx.obj['config']
-    tenant = config.get('core.tenant', None)
+    tenant = config.get('core.tenant', "")
     return tenant
 
-def get_headers(ctx=None):
-    """ Get the headers which may or may not contain a tenant """
-    headers={}
-    ctx = ctx or click.get_current_context()
-    tenant = get_tenant(ctx=ctx)
-    if tenant:
-        headers={"Cray-Tenant-Name":tenant}
-    return headers
-
 
 def hostname_to_name(hostname=None, ctx=None):
     """ Convert hostname to name value for saving as filename"""