Merge pull request #2712 from StackStorm/feature/client_api_key

Add API key support to st2client
StackStorm · May 26, 2016 · c86dc42 · c86dc42
2 parents 4bac378 + 183519a
commit c86dc42
Show file tree

Hide file tree

Showing 8 changed files with 53 additions and 4 deletions.
diff --git a/st2client/st2client/client.py b/st2client/st2client/client.py
@@ -37,7 +37,7 @@
 
 class Client(object):
     def __init__(self, base_url=None, auth_url=None, api_url=None, api_version=None, cacert=None,
-                 debug=False, token=None):
+                 debug=False, token=None, api_key=None):
         # Get CLI options. If not given, then try to get it from the environment.
         self.endpoints = dict()
 
@@ -79,6 +79,11 @@ def __init__(self, base_url=None, auth_url=None, api_url=None, api_version=None,
 
         self.token = token
 
+        if api_key:
+            os.environ['ST2_API_KEY'] = api_key
+
+        self.api_key = api_key
+
         # Instantiate resource managers and assign appropriate API endpoint.
         self.managers = dict()
         self.managers['Token'] = ResourceManager(

diff --git a/st2client/st2client/commands/resource.py b/st2client/st2client/commands/resource.py
@@ -39,6 +39,8 @@ def decorate(*args, **kwargs):
         ns = args[1]
         if getattr(ns, 'token', None):
             kwargs['token'] = ns.token
+        if getattr(ns, 'api_key', None):
+            kwargs['api_key'] = ns.api_key
         return func(*args, **kwargs)
     return decorate
 
@@ -118,6 +120,10 @@ def __init__(self, resource, *args, **kwargs):
                                      help='Access token for user authentication. '
                                           'Get ST2_AUTH_TOKEN from the environment '
                                           'variables by default.')
+            self.parser.add_argument('--api-key', dest='api_key',
+                                     help='Api Key for user authentication. '
+                                          'Get ST2_API_KEY from the environment '
+                                          'variables by default.')
 
         # Formatter flags
         self.parser.add_argument('-j', '--json',

diff --git a/st2client/st2client/config_parser.py b/st2client/st2client/config_parser.py
@@ -79,6 +79,10 @@
         'password': {
             'type': 'string',
             'default': None
+        },
+        'api_key': {
+            'type': 'string',
+            'default': None
         }
     },
     'api': {

diff --git a/st2client/st2client/models/core.py b/st2client/st2client/models/core.py
@@ -32,6 +32,8 @@ def add_auth_token_to_kwargs_from_env(func):
     def decorate(*args, **kwargs):
         if not kwargs.get('token') and os.environ.get('ST2_AUTH_TOKEN', None):
             kwargs['token'] = os.environ.get('ST2_AUTH_TOKEN')
+        if not kwargs.get('api_key') and os.environ.get('ST2_API_KEY', None):
+            kwargs['api_key'] = os.environ.get('ST2_API_KEY')
         return func(*args, **kwargs)
     return decorate
 

diff --git a/st2client/st2client/shell.py b/st2client/st2client/shell.py
@@ -79,6 +79,7 @@
     'auth_url': ['auth', 'url'],
     'api_url': ['api', 'url'],
     'api_version': ['general', 'api_version'],
+    'api_key': ['credentials', 'api_key'],
     'cacert': ['general', 'cacert'],
     'debug': ['cli', 'debug']
 }
@@ -308,7 +309,9 @@ def get_client(self, args, debug=False):
         # or as a command line argument
         env_var_token = os.environ.get('ST2_AUTH_TOKEN', None)
         cli_argument_token = getattr(args, 'token', None)
-        if env_var_token or cli_argument_token:
+        env_var_api_key = os.environ.get('ST2_API_KEY', None)
+        cli_argument_api_key = getattr(args, 'api_key', None)
+        if env_var_token or cli_argument_token or env_var_api_key or cli_argument_api_key:
             return client
 
         # If credentials are provided in the CLI config use them and try to authenticate

diff --git a/st2client/st2client/utils/httpclient.py b/st2client/st2client/utils/httpclient.py
@@ -36,11 +36,18 @@ def decorate(*args, **kwargs):
 
 def add_auth_token_to_headers(func):
     def decorate(*args, **kwargs):
+        headers = kwargs.get('headers', dict())
+
         token = kwargs.pop('token', None)
         if token:
-            headers = kwargs.get('headers', dict())
             headers['X-Auth-Token'] = str(token)
             kwargs['headers'] = headers
+
+        api_key = kwargs.pop('api_key', None)
+        if api_key:
+            headers['St2-Api-Key'] = str(api_key)
+            kwargs['headers'] = headers
+
         return func(*args, **kwargs)
     return decorate
 

diff --git a/st2client/tests/unit/test_auth.py b/st2client/tests/unit/test_auth.py
@@ -45,6 +45,7 @@ def __init__(self, *args, **kwargs):
         super(TestAuthToken, self).__init__(*args, **kwargs)
         self.parser = argparse.ArgumentParser()
         self.parser.add_argument('-t', '--token', dest='token')
+        self.parser.add_argument('--api-key', dest='api_key')
         self.shell = shell.Shell()
 
     def setUp(self):
@@ -59,6 +60,8 @@ def tearDown(self):
         # Clean up environment.
         if 'ST2_AUTH_TOKEN' in os.environ:
             del os.environ['ST2_AUTH_TOKEN']
+        if 'ST2_API_KEY' in os.environ:
+            del os.environ['ST2_API_KEY']
         if 'ST2_BASE_URL' in os.environ:
             del os.environ['ST2_BASE_URL']
 
@@ -74,12 +77,23 @@ def test_decorate_auth_token_by_cli(self):
         args = self.parser.parse_args(args=['--token', token])
         self.assertDictEqual(self._mock_run(args), {'token': token})
 
+    def test_decorate_api_key_by_cli(self):
+        token = uuid.uuid4().hex
+        args = self.parser.parse_args(args=['--api-key', token])
+        self.assertDictEqual(self._mock_run(args), {'api_key': token})
+
     def test_decorate_auth_token_by_env(self):
         token = uuid.uuid4().hex
         os.environ['ST2_AUTH_TOKEN'] = token
         args = self.parser.parse_args(args=[])
         self.assertDictEqual(self._mock_run(args), {'token': token})
 
+    def test_decorate_api_key_by_env(self):
+        token = uuid.uuid4().hex
+        os.environ['ST2_API_KEY'] = token
+        args = self.parser.parse_args(args=[])
+        self.assertDictEqual(self._mock_run(args), {'api_key': token})
+
     def test_decorate_without_auth_token(self):
         args = self.parser.parse_args(args=[])
         self.assertDictEqual(self._mock_run(args), {})
@@ -96,6 +110,13 @@ def test_decorate_auth_token_to_http_headers(self):
         self.assertIn('headers', kwargs)
         self.assertDictEqual(kwargs['headers'], expected)
 
+    def test_decorate_api_key_to_http_headers(self):
+        token = uuid.uuid4().hex
+        kwargs = self._mock_http('/', api_key=token)
+        expected = {'content-type': 'application/json', 'St2-Api-Key': token}
+        self.assertIn('headers', kwargs)
+        self.assertDictEqual(kwargs['headers'], expected)
+
     def test_decorate_without_auth_token_to_http_headers(self):
         kwargs = self._mock_http('/', auth=('stanley', 'stanley'))
         expected = {'content-type': 'application/json'}

diff --git a/st2client/tests/unit/test_config_parser.py b/st2client/tests/unit/test_config_parser.py
@@ -55,7 +55,8 @@ def test_parse(self):
             },
             'credentials': {
                 'username': 'test1',
-                'password': 'test1'
+                'password': 'test1',
+                'api_key': None
             },
             'api': {
                 'url': 'http://127.0.0.1:9101/v1'