Add support for authorization codes

Author: Kevin D Smith

swat/cas/connection.py

@@ -32,10 +32,11 @@
 import os
 import random
 import re
+import requests
 import six
 import warnings
 import weakref
-from six.moves.urllib.parse import urlparse
+from six.moves.urllib.parse import urlparse, urlencode, urljoin
 from . import rest
 from .. import clib
 from .. import config as cf
@@ -177,6 +178,8 @@ class CAS(object):
         Base path of URL when using the REST protocol.
     ssl_ca_list : string, optional
         The path to the SSL certificates for the CAS server.
+    authcode : string, optional
+        Authorization code from SASLogon used to retrieve an OAuth token.
     **kwargs : any, optional
         Arbitrary keyword arguments used for internal purposes only.
 
@@ -346,7 +349,7 @@ def _get_connection_info(cls, hostname, port, username, password, protocol, path
     def __init__(self, hostname=None, port=None, username=None, password=None,
                  session=None, locale=None, nworkers=None, name=None,
                  authinfo=None, protocol=None, path=None, ssl_ca_list=None,
-                 auth_code=None, **kwargs):
+                 authcode=None, **kwargs):
 
         # Filter session options allowed as parameters
         _kwargs = {}
@@ -392,11 +395,11 @@ def __init__(self, hostname=None, port=None, username=None, password=None,
             soptions = a2n(getsoptions(session=session, locale=locale,
                                        nworkers=nworkers, protocol=protocol))
 
-            # Check for auth_code authentication
-            auth_code = auth_code or cf.get_option('cas.auth_code')
-            if protocol in ['http', 'https'] and auth_code:
+            # Check for authcode authentication
+            authcode = authcode or cf.get_option('cas.authcode')
+            if protocol in ['http', 'https'] and authcode:
                 username = None
-                password = self.get_token(auth_code=auth_code, url=hostname)
+                password = type(self)._get_token(authcode=authcode, url=hostname)
 
         # Create error handler
         try:
@@ -529,33 +532,41 @@ def _id_generator():
                 num = num + 1
         self._id_generator = _id_generator()
 
-    def _get_token(self, username=None, password=None, auth_code=None,
-                   client_id=None, client_secret=None, base_url=None):
+    @classmethod
+    def _get_token(cls, username=None, password=None, authcode=None,
+                   client_id=None, client_secret=None, url=None):
         ''' Retrieve token from Viya installation '''
-        headers = {'Accept': 'application/vnd.sas.compute.session+json',
-                   'Content-Type': 'application/x-www-form-urlencoded'}
+        from .rest.connection import _print_request, _setup_ssl
+
+        with requests.Session() as req_sess:
+
+            _setup_ssl(req_sess)
+
+            req_sess.headers.update({'Accept': 'application/vnd.sas.compute.session+json',
+                                     'Content-Type': 'application/x-www-form-urlencoded'})
 
-        auth_code = auth_code or cf.get_option('cas.auth_code')
-        if auth_code:
             client_id = client_id or cf.get_option('cas.client_id') or 'SWAT'
-            client_secret = client_secret or cf.get_option('cas.client_secret') or ''
-            body = {'grant_type': 'authorization_code', 'code': auth_code,
-                    'client_id': client_id, 'client_secret': client_secret}
-        else:
-            user = client_id or cf.get_option('cas.username')
-            password = client_secret or cf.get_option('cas.token')
-            body = {'grant_type': 'password', 'username': user,
-                    'password': password}
-
-        resp = requests.post(base_url + '/SASLogon/oauth/token',
-                             headers=headers, user=('sas.tkmtrb', ''),
-                             data=urlparse.quote_plus(body))
-
-        if resp.status_code >= 300:
-            raise SWATError('Token request resulted in a status of %s' %
-                            resp.status_code)
- 
-        return resp['access_token']
+
+            authcode = authcode or cf.get_option('cas.authcode')
+            if authcode:
+                client_secret = client_secret or cf.get_option('cas.client_secret') or ''
+                body = {'grant_type': 'authorization_code', 'code': authcode,
+                        'client_id': client_id, 'client_secret': client_secret}
+            else:
+                username = username or cf.get_option('cas.username')
+                password = password or cf.get_option('cas.token')
+                body = {'grant_type': 'password', 'username': username,
+                        'password': password}
+
+            resp = req_sess.post(urljoin(url, '/SASLogon/oauth/token'),
+                                 auth=(client_id, ''),
+                                 data=urlencode(body))
+
+            if resp.status_code >= 300:
+                raise SWATError('Token request resulted in a status of %s' %
+                                resp.status_code)
+
+            return resp.json()['access_token']
 
     def _gen_id(self):
         ''' Generate an ID unique to the session '''

swat/cas/rest/connection.py

@@ -175,6 +175,24 @@ def init_poolmanager(self, connections, maxsize,
                                                                **pool_kwargs)
 
 
+def _setup_ssl(req_sess):
+    ''' Configure a Requests session for SSL '''
+    if os.environ.get('SSLREQCERT', 'y').lower() in ['n', 'no', '0',
+                                                     'f', 'false', 'off']:
+        req_sess.verify = False
+    elif 'CAS_CLIENT_SSL_CA_LIST' in os.environ:
+        req_sess.verify = os.path.expanduser(
+            os.environ['CAS_CLIENT_SSL_CA_LIST'])
+    elif 'SAS_TRUSTED_CA_CERTIFICATES_PEM_FILE' in os.environ:
+        req_sess.verify = os.path.expanduser(
+            os.environ['SAS_TRUSTED_CA_CERTIFICATES_PEM_FILE'])
+    elif 'SSLCALISTLOC' in os.environ:
+        req_sess.verify = os.path.expanduser(
+            os.environ['SSLCALISTLOC'])
+    elif 'REQUESTS_CA_BUNDLE' not in os.environ:
+        req_sess.mount('https://', SSLContextAdapter())
+
+
 class REST_CASConnection(object):
     '''
     Create a REST CAS connection
@@ -286,22 +304,10 @@ def __init__(self, hostname, port, username, password, soptions, error):
 
         self._req_sess = requests.Session()
 
-        if os.environ.get('SSLREQCERT', 'y').lower() in ['n', 'no', '0',
-                                                         'f', 'false', 'off']:
-            self._req_sess.verify = False
-        elif 'CAS_CLIENT_SSL_CA_LIST' in os.environ:
-            self._req_sess.verify = os.path.expanduser(
-                os.environ['CAS_CLIENT_SSL_CA_LIST'])
-        elif 'SAS_TRUSTED_CA_CERTIFICATES_PEM_FILE' in os.environ:
-            self._req_sess.verify = os.path.expanduser(
-                os.environ['SAS_TRUSTED_CA_CERTIFICATES_PEM_FILE'])
-        elif 'SSLCALISTLOC' in os.environ:
-            self._req_sess.verify = os.path.expanduser(
-                os.environ['SSLCALISTLOC'])
-        elif 'REQUESTS_CA_BUNDLE' not in os.environ:
-            self._req_sess.mount('https://', SSLContextAdapter())
+        _setup_ssl(self._req_sess)
 
         self._req_sess.headers.update({
+            'Accept': 'application/json',
             'Content-Type': 'application/json',
             'Content-Length': '0',
             'Authorization': self._auth,
@@ -491,6 +497,7 @@ def invoke(self, action_name, kwargs):
 
         post_data = a2u(kwargs).encode('utf-8')
         self._req_sess.headers.update({
+            'Accept': 'application/json',
             'Content-Type': 'application/json',
             'Content-Length': str(len(post_data)),
         })
@@ -706,6 +713,7 @@ def upload(self, file_name, params):
             data = datafile.read()
 
         self._req_sess.headers.update({
+            'Accept': 'application/json',
             'Content-Type': 'application/octet-stream',
             'Content-Length': str(len(data)),
             'JSON-Parameters': json.dumps(_normalize_params(params))

swat/config.py

@@ -173,20 +173,20 @@ def check_tz(value):
                 'Sets the port number for the CAS server.',
                 environ='CAS_PORT')
 
-register_option('cas.auth_code', 'string', check_string, None,
-                'Sets the auth_code for retrieving an OAuth token from '
+register_option('cas.authcode', 'string', check_string, None,
+                'Sets the authorization code for retrieving an OAuth token from '
                 'a Viya deployment.',
-                environ='CAS_AUTHCODE')
+                environ=['CAS_AUTHCODE', 'VIYA_AUTHCODE'])
 
 register_option('cas.client_id', 'string', check_string, 'SWAT',
                 'Sets the client ID for retrieving an OAuth token from '
                 'a Viya deployment.',
-                environ='CAS_CLIENT_ID')
+                environ=['CAS_CLIENT_ID', 'VIYA_CLIENT_ID'])
 
 register_option('cas.client_secret', 'string', check_string, '',
                 'Sets the client secret for retrieving an OAuth token from '
                 'a Viya deployment.',
-                environ='CAS_CLIENT_SECRET')
+                environ=['CAS_CLIENT_SECRET', 'VIYA_CLIENT_SECRET'])
 
 register_option('cas.protocol', 'string',
                 functools.partial(check_string,