88import copy
99import json
1010import logging
11- import platform
1211import tempfile
1312import time
1413import uuid
1817from threading import Lock , Thread
1918
2019from .auth_keypair import AuthByKeyPair
21- from .compat import IS_LINUX , TO_UNICODE , urlencode
20+ from .compat import IS_LINUX , IS_WINDOWS , IS_MACOS , TO_UNICODE , urlencode
2221from .constants import (
2322 HTTP_HEADER_ACCEPT ,
2423 HTTP_HEADER_CONTENT_TYPE ,
2524 HTTP_HEADER_SERVICE_NAME ,
2625 HTTP_HEADER_USER_AGENT ,
2726 PARAMETER_CLIENT_STORE_TEMPORARY_CREDENTIAL ,
28- PARAMETER_CLIENT_USE_SECURE_STORAGE_FOR_TEMPORARY_CREDENTIAL ,
2927)
3028from .description import COMPILER , IMPLEMENTATION , OPERATING_SYSTEM , PLATFORM , PYTHON_VERSION
3129from .errorcode import ER_FAILED_TO_CONNECT_TO_DB
32- from .errors import BadGatewayError , DatabaseError , Error , ForbiddenError , ServiceUnavailableError
30+ from .errors import (
31+ BadGatewayError ,
32+ DatabaseError ,
33+ Error ,
34+ ForbiddenError ,
35+ ServiceUnavailableError ,
36+ MissingDependencyError ,
37+ )
3338from .network import (
3439 ACCEPT_TYPE_APPLICATION_SNOWFLAKE ,
3540 CONTENT_TYPE_APPLICATION_JSON ,
4146
4247logger = logging .getLogger (__name__ )
4348
49+ try :
50+ import keyring
51+ except ImportError as ie :
52+ keyring = None
53+ logger .debug ('Failed to import keyring module. err=[%s]' , ie )
54+
4455# Cache directory
4556CACHE_ROOT_DIR = getenv ('SF_TEMPORARY_CREDENTIAL_CACHE_DIR' ) or \
4657 expanduser ("~" ) or tempfile .gettempdir ()
47- if platform . system () == 'Windows' :
58+ if IS_WINDOWS :
4859 CACHE_DIR = path .join (CACHE_ROOT_DIR , 'AppData' , 'Local' , 'Snowflake' ,
4960 'Caches' )
50- elif platform . system () == 'Darwin' :
61+ elif IS_MACOS :
5162 CACHE_DIR = path .join (CACHE_ROOT_DIR , 'Library' , 'Caches' , 'Snowflake' )
5263else :
5364 CACHE_DIR = path .join (CACHE_ROOT_DIR , '.cache' , 'snowflake' )
7788# keyring
7889KEYRING_SERVICE_NAME = "net.snowflake.temporary_token"
7990KEYRING_USER = "temp_token"
91+ KEYRING_DRIVER_NAME = "SNOWFLAKE-PYTHON-DRIVER"
8092
8193
8294class Auth (object ):
@@ -91,22 +103,28 @@ def __init__(self, rest):
91103 def base_auth_data (user , account , application ,
92104 internal_application_name ,
93105 internal_application_version ,
94- ocsp_mode ):
106+ ocsp_mode , login_timeout ,
107+ network_timeout = None ,
108+ store_temp_cred = None ):
95109 return {
96- u'data' : {
97- u"CLIENT_APP_ID" : internal_application_name ,
98- u"CLIENT_APP_VERSION" : internal_application_version ,
99- u"SVN_REVISION" : VERSION [3 ],
100- u"ACCOUNT_NAME" : account ,
101- u"LOGIN_NAME" : user ,
102- u"CLIENT_ENVIRONMENT" : {
103- u"APPLICATION" : application ,
104- u"OS" : OPERATING_SYSTEM ,
105- u"OS_VERSION" : PLATFORM ,
106- u"PYTHON_VERSION" : PYTHON_VERSION ,
107- u"PYTHON_RUNTIME" : IMPLEMENTATION ,
108- u"PYTHON_COMPILER" : COMPILER ,
109- u"OCSP_MODE" : ocsp_mode .name ,
110+ 'data' : {
111+ "CLIENT_APP_ID" : internal_application_name ,
112+ "CLIENT_APP_VERSION" : internal_application_version ,
113+ "SVN_REVISION" : VERSION [3 ],
114+ "ACCOUNT_NAME" : account ,
115+ "LOGIN_NAME" : user ,
116+ "CLIENT_ENVIRONMENT" : {
117+ "APPLICATION" : application ,
118+ "OS" : OPERATING_SYSTEM ,
119+ "OS_VERSION" : PLATFORM ,
120+ "PYTHON_VERSION" : PYTHON_VERSION ,
121+ "PYTHON_RUNTIME" : IMPLEMENTATION ,
122+ "PYTHON_COMPILER" : COMPILER ,
123+ "OCSP_MODE" : ocsp_mode .name ,
124+ "TRACING" : logger .getEffectiveLevel (),
125+ "LOGIN_TIMEOUT" : login_timeout ,
126+ "NETWORK_TIMEOUT" : network_timeout ,
127+ "CLIENT_STORE_TEMPORARY_CREDENTIAL" : store_temp_cred ,
110128 }
111129 },
112130 }
@@ -132,11 +150,22 @@ def authenticate(self, auth_instance, account, user,
132150 headers [HTTP_HEADER_SERVICE_NAME ] = \
133151 session_parameters [HTTP_HEADER_SERVICE_NAME ]
134152 url = u"/session/v1/login-request"
153+ if session_parameters is not None \
154+ and PARAMETER_CLIENT_STORE_TEMPORARY_CREDENTIAL in session_parameters :
155+ store_temp_cred = session_parameters [
156+ PARAMETER_CLIENT_STORE_TEMPORARY_CREDENTIAL ]
157+ else :
158+ store_temp_cred = None
159+
135160 body_template = Auth .base_auth_data (
136161 user , account , self ._rest ._connection .application ,
137162 self ._rest ._connection ._internal_application_name ,
138163 self ._rest ._connection ._internal_application_version ,
139- self ._rest ._connection ._ocsp_mode ())
164+ self ._rest ._connection ._ocsp_mode (),
165+ self ._rest ._connection ._login_timeout ,
166+ self ._rest ._connection ._network_timeout ,
167+ store_temp_cred ,
168+ )
140169
141170 body = copy .deepcopy (body_template )
142171 # updating request body
@@ -317,10 +346,10 @@ def post_request_wrapper(self, url, headers, body):
317346 id_token = ret [u'data' ].get (u'idToken' )
318347 )
319348 if self ._rest ._connection .consent_cache_id_token :
320- write_temporary_credential_file (
321- account , user , self ._rest .id_token ,
349+ write_temporary_credential (
350+ self . _rest . _host , account , user , self ._rest .id_token ,
322351 session_parameters .get (
323- PARAMETER_CLIENT_USE_SECURE_STORAGE_FOR_TEMPORARY_CREDENTIAL ))
352+ PARAMETER_CLIENT_STORE_TEMPORARY_CREDENTIAL ))
324353 if u'sessionId' in ret [u'data' ]:
325354 self ._rest ._connection ._session_id = ret [u'data' ][u'sessionId' ]
326355 if u'sessionInfo' in ret [u'data' ]:
@@ -333,17 +362,26 @@ def post_request_wrapper(self, url, headers, body):
333362
334363 return session_parameters
335364
336- def read_temporary_credential (self , account , user , session_parameters ):
337- if session_parameters .get (PARAMETER_CLIENT_STORE_TEMPORARY_CREDENTIAL ):
338- read_temporary_credential_file (
339- session_parameters .get (
340- PARAMETER_CLIENT_USE_SECURE_STORAGE_FOR_TEMPORARY_CREDENTIAL )
341- )
342- id_token = TEMPORARY_CREDENTIAL .get (
343- account .upper (), {}).get (user .upper ())
365+ def read_temporary_credential (self , host , account , user , session_parameters ):
366+ if session_parameters .get (PARAMETER_CLIENT_STORE_TEMPORARY_CREDENTIAL , False ):
367+ id_token = None
368+ if IS_MACOS or IS_WINDOWS :
369+ if not keyring :
370+ # we will leave the exception for write_temporary_credential function to raise
371+ return False
372+ new_target = convert_target (host , user )
373+ try :
374+ id_token = keyring .get_password (new_target , user .upper ())
375+ except keyring .errors .KeyringError as ke :
376+ logger .debug ("Could not retrieve id_token from secure storage : {}" .format (str (ke )))
377+ elif IS_LINUX :
378+ read_temporary_credential_file ()
379+ id_token = TEMPORARY_CREDENTIAL .get (
380+ account .upper (), {}).get (user .upper ())
381+ else :
382+ logger .debug ("connection parameter enable_sso_temporary_credential not set or OS not support" )
344383 if id_token :
345384 self ._rest .id_token = id_token
346- if self ._rest .id_token :
347385 try :
348386 self ._rest ._id_token_session ()
349387 return True
@@ -354,11 +392,31 @@ def read_temporary_credential(self, account, user, session_parameters):
354392 return False
355393
356394
357- def write_temporary_credential_file (
358- account , user , id_token ,
359- use_secure_storage_for_temporary_credential = False ):
360- if not CACHE_DIR or not id_token :
361- # no cache is enabled or no id_token is given
395+ def write_temporary_credential (host , account , user , id_token , store_temporary_credential = False ):
396+ if not id_token :
397+ logger .debug ("no ID token is given when try to store temporary credential" )
398+ return
399+ if IS_MACOS or IS_WINDOWS :
400+ if not keyring :
401+ raise MissingDependencyError ("Please install keyring module to enable SSO token cache feature." )
402+
403+ new_target = convert_target (host , user )
404+ try :
405+ keyring .set_password (new_target , user .upper (), id_token )
406+ except keyring .errors .KeyringError as ke :
407+ logger .debug ("Could not store id_token to keyring, %s" , str (ke ))
408+ elif IS_LINUX and store_temporary_credential :
409+ write_temporary_credential_file (host , account , user , id_token )
410+ else :
411+ logger .debug ("connection parameter client_store_temporary_credential not set or OS not support" )
412+
413+
414+ def write_temporary_credential_file (host , account , user , id_token ):
415+ """
416+ Write temporary credential file when OS is Linux
417+ """
418+ if not CACHE_DIR :
419+ # no cache is enabled
362420 return
363421 global TEMPORARY_CREDENTIAL
364422 global TEMPORARY_CREDENTIAL_LOCK
@@ -377,27 +435,19 @@ def write_temporary_credential_file(
377435 "write the temporary credential file: %s" ,
378436 TEMPORARY_CREDENTIAL_FILE )
379437 try :
380- if IS_LINUX or not use_secure_storage_for_temporary_credential :
381- with codecs .open (TEMPORARY_CREDENTIAL_FILE , 'w' ,
382- encoding = 'utf-8' , errors = 'ignore' ) as f :
383- json .dump (TEMPORARY_CREDENTIAL , f )
384- else :
385- import keyring
386- keyring .set_password (
387- KEYRING_SERVICE_NAME , KEYRING_USER ,
388- json .dumps (TEMPORARY_CREDENTIAL ))
389-
438+ with codecs .open (TEMPORARY_CREDENTIAL_FILE , 'w' ,
439+ encoding = 'utf-8' , errors = 'ignore' ) as f :
440+ json .dump (TEMPORARY_CREDENTIAL , f )
390441 except Exception as ex :
391442 logger .debug ("Failed to write a credential file: "
392443 "file=[%s], err=[%s]" , TEMPORARY_CREDENTIAL_FILE , ex )
393444 finally :
394445 unlock_temporary_credential_file ()
395446
396447
397- def read_temporary_credential_file (
398- use_secure_storage_for_temporary_credential = False ):
448+ def read_temporary_credential_file ():
399449 """
400- Read temporary credential file
450+ Read temporary credential file when OS is Linux
401451 """
402452 if not CACHE_DIR :
403453 # no cache is enabled
@@ -416,15 +466,9 @@ def read_temporary_credential_file(
416466 "write the temporary credential file: %s" ,
417467 TEMPORARY_CREDENTIAL_FILE )
418468 try :
419- if IS_LINUX or not use_secure_storage_for_temporary_credential :
420- with codecs .open (TEMPORARY_CREDENTIAL_FILE , 'r' ,
421- encoding = 'utf-8' , errors = 'ignore' ) as f :
422- TEMPORARY_CREDENTIAL = json .load (f )
423- else :
424- import keyring
425- f = keyring .get_password (
426- KEYRING_SERVICE_NAME , KEYRING_USER ) or "{}"
427- TEMPORARY_CREDENTIAL = json .loads (f )
469+ with codecs .open (TEMPORARY_CREDENTIAL_FILE , 'r' ,
470+ encoding = 'utf-8' , errors = 'ignore' ) as f :
471+ TEMPORARY_CREDENTIAL = json .load (f )
428472 return TEMPORARY_CREDENTIAL
429473 except Exception as ex :
430474 logger .debug ("Failed to read a credential file. The file may not"
@@ -456,26 +500,34 @@ def unlock_temporary_credential_file():
456500 return False
457501
458502
459- def delete_temporary_credential_file (
460- use_secure_storage_for_temporary_credential = False ):
461- """
462- Delete temporary credential file and its lock file
463- """
464- global TEMPORARY_CREDENTIAL_FILE
465- if IS_LINUX or not use_secure_storage_for_temporary_credential :
503+ def delete_temporary_credential (host , user , store_temporary_credential = False ):
504+ if (IS_MACOS or IS_WINDOWS ) and keyring :
505+ new_target = convert_target (host , user )
466506 try :
467- remove (TEMPORARY_CREDENTIAL_FILE )
468- except Exception as ex :
469- logger .debug ("Failed to delete a credential file: "
470- "file=[%s], err=[%s]" , TEMPORARY_CREDENTIAL_FILE , ex )
471- else :
472- try :
473- import keyring
474- keyring .delete_password (KEYRING_SERVICE_NAME , KEYRING_USER )
507+ keyring .delete_password (new_target , user .upper ())
475508 except Exception as ex :
476509 logger .debug ("Failed to delete credential in the keyring: err=[%s]" ,
477510 ex )
511+ elif IS_LINUX and store_temporary_credential :
512+ delete_temporary_credential_file ()
513+
514+
515+ def delete_temporary_credential_file ():
516+ """
517+ Delete temporary credential file and its lock file
518+ """
519+ global TEMPORARY_CREDENTIAL_FILE
520+ try :
521+ remove (TEMPORARY_CREDENTIAL_FILE )
522+ except Exception as ex :
523+ logger .debug ("Failed to delete a credential file: "
524+ "file=[%s], err=[%s]" , TEMPORARY_CREDENTIAL_FILE , ex )
478525 try :
479526 removedirs (TEMPORARY_CREDENTIAL_FILE_LOCK )
480527 except Exception as ex :
481528 logger .debug ("Failed to delete credential lock file: err=[%s]" , ex )
529+
530+
531+ def convert_target (host , user ):
532+ return "{host}:{user}:{driver}" .format (
533+ host = host .upper (), user = user .upper (), driver = KEYRING_DRIVER_NAME )
0 commit comments