o
    RDi                     @  s  d dl mZ d dlZd dlZd dlZd dlZd dlZd dlZd dlZd dl	Z	d dl
m
Z
 d dlmZ d dlmZ d dlmZmZ d dlmZ d dlZd dlZdd	lmZmZ dd
lmZ ddlmZmZ ddlmZ ddl m!Z! ee"Z#erzd dl$Z$dddZ%dddZ&G dd dZ'dS )    )annotationsN)datetime)	getLogger)Path)AnyAnyStr)
getproxies   )
IS_WINDOWSurlparse)SnowflakeCursor)SessionManagerSessionManagerFactory)&extract_top_level_domain_from_hostname)urllib3ddict[str, dict[str, Any]]c                 C  s`   i }|   D ]'\}}t|tr| }t|tr| }n	t|tr&t|}|||i q|S N)items
isinstancebytesdecodedict_decode_dictupdate)r   resultkeyvalue r   b/var/www/Datamplify/venv/lib/python3.10/site-packages/snowflake/connector/connection_diagnostic.pyr       s   



r   	allowlistlist[dict[str, Any]]c                 C  sF   t | tr!tdd | D r!zt|  W dS  ty    Y dS w dS )Nc                 s  s    | ]}t |tV  qd S r   )r   r   ).0itemr   r   r   	<genexpr>.       

z+_is_list_of_json_objects.<locals>.<genexpr>TF)r   listalljsondumps	TypeError)r    r   r   r   _is_list_of_json_objects-   s   
r+   c                   @  s   e Zd ZdZ							d?d@ddZdAddZ			dBdCddZdDd d!ZdEd#d$ZdFdGd%d&Z		dHdId'd(Z
dJd+d,ZdKd.d/ZdDd0d1ZdDd2d3ZdDd4d5ZdLd9d:ZdMd=d>ZdS )NConnectionDiagnosticzyImplementation of a connection test utility for Snowflake connector

    Use new ConnectionTest() to get the object.
    Naccountstrhostconnection_diag_log_path
str | Noneconnection_diag_allowlist_path
proxy_host
proxy_port
proxy_userproxy_passwordsession_managerSessionManager | NonereturnNonec
                 C  s  || _ || _g g g g g g g d| _d}
| |
d| j   | |
d| j  t|}d| d| | jv rb|d| ddd	 d|  | _td
| d| d | |
d| d| j  n|| _g | _g | _	i | _
d| _|| _|| _|| _|| _| jd u rtd}d| _nt d }d| _| |\| _| _| _| _| | j |d urt|nd | _|d urt|nd | _t | _| jd u rt| j| _n+| j std| j d t| j| _n| j std| j d t| j| _| jd | _td| j  | jd ur>| j s&td| j d td d | _n| j s>td| j d td d | _d| _|   rT| j!d| j  d | _n	| j!d!|  d"| _"d | _#|	ro|	j$d"d#| _'d S t%j&d"d#| _'d S )$N)INITIALPROXYSNOWFLAKE_URLSTAGEOCSP_RESPONDEROUT_OF_BAND_TELEMETRYIGNOREr;   zSpecified snowflake account: z!Host based on specified account: .z.snowflakecomputing.z.snowr	   r   z+Account should not have snowflakecomputing.z in it. You provided z.  Continuing with fixed host.z%We removed extra .snowflakecomputing.z and will continue with host: paramsHTTPS_PROXYenvironmenthttpssystemzPath z% for connection test is not absolute.z$ for connection test does not exist.z!SnowflakeConnectionTestReport.txtzReporting to file zPath 'z0' for connection test allowlist is not absolute.zWill connect to Snowflake for allowlist json instead.  If you did not provide a valid password, please make sure to update and run again.zFile 'z/' for connection test allowlist does not exist.zRselect /* snowflake-connector-python:connection_diagnostics */ system$allowlist();zocsp.z&select system$allowlist_privatelink();zocsp.snowflakecomputing.F)use_pooling)(r-   r/   test_results%_ConnectionDiagnostic__append_messager   splitloggerwarning	ocsp_urlscrl_urls	cert_info
proxy_typer3   r4   r5   r6   osgetenvr   "_ConnectionDiagnostic__parse_proxy(_ConnectionDiagnostic__https_host_reportr   full_connection_diag_log_path#full_connection_diag_allowlist_pathtempfile
gettempdirtmpdiris_absoluteexistsreport_fileinfoallowlist_sql%_ConnectionDiagnostic__is_privatelinkappendallowlist_retrieval_successcursorcloner   get_manager_session_manager)selfr-   r/   r0   r2   r3   r4   r5   r6   r7   	host_typetop_level_domain	proxy_urlr   r   r   __init__?   s   	 








zConnectionDiagnostic.__init__rj   tuple[str, str, str, str]c                 C  s,   t |}|j}|j}|j}|j}||||fS r   )r   hostnameportusernamepassword)rg   rj   parsedr3   r4   r5   r6   r   r   r   __parse_proxy   s   z"ConnectionDiagnostic.__parse_proxy  
   r=   rn   inttimeoutrh   c                 C  s*  zH| j ||d d}ttjtj}|| | jd urL| j| jf}| jd urF| j d| j }|	d}t
| d}	d|	 d}|| n	||t|f |dkr| jd urd	| d| d
| }
|
 d| d}
|t	|
 |dd ttj}|t  ttdd}|rt|dr| j|O  _|j||d}t|d}d| d| d| d}z
|t	| W n  ty } z|  || d| d|  W Y d }~nd }~ww |!  |W S | jd urd	| d| d
| d}
nd| d}
|t	|
 |dd}|!  |d ur;d}t"#||s;|  || d| d|  W dS |  || d| d W dS  tj$ys } zdt|v rh|  || d| d|  W Y d }~dS d }~w ty } z|  || d| d|  W Y d }~dS d }~ww ) N)rh    :utf-8
zProxy-Authorization: Basic z
rs   zCONNECT z HTTP/1.1
zHost: z

i   VERIFY_X509_PARTIAL_CHAINr   verify_flags)server_hostnameTzGET / z5 HTTP/1.1

                                   Host: z

                                   User-Agent: snowflake-connector-python-diagnostic
                                   

z(: URL Check: Failed: Unknown Exception: zGET / HTTP/1.1
Host: z
Connection: close

z(200|301|cloudfront)z: URL Check: Failed: FAILED#: URL Check: Connected SuccessfullySUCCESSWRONG_VERSION_NUMBERz': URL Check: Failed: Proxy Auth Error: )%_ConnectionDiagnostic__list_ipssocketAF_INETSOCK_STREAM
settimeoutr3   r4   r5   r6   encodebase64	b64encoder   stripconnectru   sendr.   recvssl
SSLContextPROTOCOL_TLS_CLIENTload_verify_locationscertifiwheregetattrhasattrr|   wrap_socketDER_cert_to_PEM_certgetpeercert	ExceptionrJ   closeresearchSSLError)rg   r/   rn   rv   rh   connect_credsconn
proxy_addr
proxy_authcredentialsr   context_partial_flagsockcertificatehttp_requesteresponsegood_responsesr   r   r   __test_socket_get_cert   s   






z+ConnectionDiagnostic.__test_socket_get_certc                 C  sx  d }| j d u rC| jd urBz| jj| jdd d d }tt|}d| _W nK t	yA } zt
d|  W Y d }~n6d }~ww n0t| j }z
t|}d| _W n  t	yr } z| dd| d| j  d W Y d }~nd }~ww t|r|D ]8}|d	 }|d
 }|d }|dv r|| jvr| j|||d qy|dv rz
| j|||d W qy t	y   Y qyw qyd S | dd d S )NT)_is_internalr   z*Unable to do allowlist checks: exception: r;   zAllowlist was not valid json: 'zB'.  Please run 'select system$allowlist();' and validate the file z is correct.typer/   rn   r?   rn   rh   )r>   r@   zAllowlist is not a valid list of json objects. Please run 'select system$allowlist();' and provide as a json file using the connection_diag_allowlist_path option.)rW   rc   executer_   fetchallr(   loadsr.   rb   r   rL   rM   openloadrJ   r+   rN   +_ConnectionDiagnostic__test_socket_get_certrU   )rg   resultsr   results_filer   rh   r/   	host_portr   r   r   run_post_test1  sn   








z"ConnectionDiagnostic.run_post_testboolc                 C  s
   d| j v S )Nprivatelink)r/   )rg   r   r   r   __is_privatelinka  s   
z%ConnectionDiagnostic.__is_privatelinkc              
   C  s   zYt |d }| d}d|v rL|D ]5}t|jr,|  s+| || d| d q|  r=| || d| d q| || d|  qW d S | || d|  W d S  tyt } zt	d	|  W Y d }~d S d }~ww )
N   z: nslookup resultssnowflakecomputingz: private ip: z<: WARNING: this is not typical for a non-privatelink accountz: public ip: z7: WARNING: privatelink accounts must have a private ip.: z)Connectivity Test Exception in list_ips: )
r   gethostbyname_ex	ipaddress
ip_address
is_privater`   rJ   r   rL   rM   )rg   r/   rh   ipsbase_messageipr   r   r   r   
__list_ipsd  s4   
zConnectionDiagnostic.__list_ipsc                   sN  z| j |||d}d|v rtjtjj| t   t    	  
 tt  ddtt  ddd}|| j|<  fddt  D }i }|D ]}t||| d< q]|d\}	}	}
|
t|d	 v r| || d
| d n%d|v r|
t|d v r| || d
| d n| || d
| d | || d t|d	 }| || d|  t|d }| || d|  | || d|d   | || d|d   | || d|d   | || d|d   |dkr[d|v r0td|d }|D ]}| j|dd  q n| dd  d!|v rUtd|d! }|D ]}| j|dd  qEn| d"d# d|v rm| || d$|d   | || d%| j  | || d&| j  W d S W d S  ty } zt d'|  W Y d }~d S d }~ww )(Nr   zBEGIN CERTIFICATEry   z%Y%m%d%H%M%SZ)subjectissuerserialNumberversion	notBeforenotAfterc                 3  s    | ]}  |V  qd S r   )get_extension)r"   ix509r   r   r$     r%   z;ConnectionDiagnostic.__https_host_report.<locals>.<genexpr>rB   r   rx   r   subjectAltNamezK: URL Check: Failed: Certificate mismatch: Host not in subject or alt namesz: Cert info:z: subject: r   z
: issuer: z: serialNumber: r   z: version: r   z: notBefore: r   z: notAfter: r   r=   authorityInfoAccessz(https?://\S+)/r   r;   z(Unable to find ocsp URLs in certificate.crlDistributionPointsrA   z'Unable to find crl URLs in certificate.z: subjectAltName: z: crlUrls: z: ocspURLs: z2Connectivity Test Exception in https_host_report: )!r   OpenSSLcryptoload_certificateFILETYPE_PEMr   get_subjectget_components
get_issuerget_serial_numberget_versionr   strptimer.   get_notBeforer   get_notAfterrP   rangeget_extension_countget_short_name	partitionrJ   r   r   findallrN   ra   rK   rO   r   rL   rM   )rg   r/   rn   rh   r   r   
extensionsextension_datar   _host_suffixsubject_str
issuer_strocsp_urls_origurlcrl_urls_origr   r   r   __https_host_report  s   







]z(ConnectionDiagnostic.__https_host_reportr   dict[bytes, bytes]c                 C  s,   t ddtt|dddd}|S )Nz[{}"]rw   r   =,;)r   subr(   r)   r   replace)rg   r   r   r   r   r   __get_issuer_string  s
   z(ConnectionDiagnostic.__get_issuer_stringmessagec                 C  s   | j | | d|  d S )Nr   )rI   ra   )rg   rh   r   r   r   r   __append_message  s   z%ConnectionDiagnostic.__append_messagec              
   C  s<  i }d}g }|D ]}|t j v r"t j|||< t j|= || qd}t }| |d|  d| v rk| t d \}}}	}
|	d urVd|	 d|
 d| d| }nd| d| }| |d| d	| d
 |D ]	}|| t j|< qm| |dt   d}dd| d}| j	| j
 v r| | j
| j	 d }t||s| |d| j	 d| d|  d}| j|ddd || j
 v r| | j
| d }t||s| |d| d| d|  trd}ddg}|D ]	}| ||| qzWt  ddd}| jd ur3| jd ur"d| j d| j d| j d| j }n
d| j d| j }||d|d< | jj	 d)d!di|}d"|jv rN| |d# W d S W d S  ty } z@d$t|v rj| |d%|  n!d&t|v r{| |d'|  n| |d(|  W Y d }~d S W Y d }~d S W Y d }~d S d }~ww )*N)
HTTP_PROXYrD   https_proxy
http_proxyr<   z/Proxies with Env vars removed(SYSTEM PROXIES): rF   http://rx   @zIf there are failures, try using the SYSTEM PROXY: On Windows, do
                                                 "set HTTPS_PROXY='z]'".  On Linux/Mac, do
                                                  "export HTTPS_PROXY='z'" z-Proxies with Env vars restored(ENV PROXIES): )z!C=US; O=Google Trust Services LLCzC=US; O=AmazonzC=US; O=DigiCert Incz(^z|^)r   z/There is likely a proxy because the issuer for z is not correct. Got z and expected one of zwww.google.comrs   rA   r   z)Software\Microsoft\Windows\CurrentVersionHKEY_CURRENT_USERHKEY_LOCAL_MACHINErt   F)rv   verify)httprF   proxies!https://nonexistentdomain.invalidrH   zdoes not existz5It is likely there is a proxy based on HTTP response.NewConnectionErrorzfProxy check using invalid URL did not show proxy: Review result, but you can probably ignore: Result: 
ProxyErrorz2It is likely there is a proxy based on Exception: zJCould not determine if a proxy does or does not exist based on Exception: )r  )rR   environkeysgetra   r   rJ   rT   joinr/   rP   (_ConnectionDiagnostic__get_issuer_stringr   r   rU   r
   (_ConnectionDiagnostic__walk_win_registryr   disable_warningsr3   r5   r6   r4   rf   textr   r.   )rg   env_proxy_backup
proxy_keysrestore_keys	proxy_keyrh   system_proxiesr3   r4   r5   r6   proxy_url_examplerestore_keycert_authoritiescheck_patternr   	test_hostregistry_start_keyhkey_stringshkey_strrequest_kwargsrj   respr   r   r   r   __check_for_proxies  s   

&

z(ConnectionDiagnostic.__check_for_proxiesc                 C  s8   |    tt| j| _| jD ]
}| j|ddd qd S )NP   r?   r   )(_ConnectionDiagnostic__check_for_proxiesr&   setrN   r   )rg   r   r   r   r   run_testq  s
   
zConnectionDiagnostic.run_testc                 C  s  d}d | jd }| d| d}d | jd }| d| d}d | jd }| d| d}| jrFd | jd }| d	| d}n| d
}| d}d | jd }| jra| d}n| d}| d| d}| jrd | jd }| d| d}t| | j| d S )NzG=========Connectivity diagnostic report================================rz   r;   r<   zI
=========Proxy information - These are best guesses, not guarantees====
r=   zI
=========Snowflake URL information=====================================
r>   zt
=========Snowflake Stage information===================================
We retrieved stage info from the allowlist
z
=========Snowflake Stage information - Unavailable=====================
We could not connect to Snowflake to get allowlist, so we do not have stage
diagnostic info
zH
=========Snowflake OCSP information====================================r?   zi
We were able to retrieve system allowlist.
These OCSP hosts came from the certificate and the allowlist.z^
We were unable to retrieve system allowlist.
These OCSP hosts only came from the certificate.r@   zI
=========Snowflake Out of bound telemetry check========================
)r
  rI   rb   rL   debugr]   
write_text)rg   r   initial_joined_resultsproxy_joined_resultssnowflake_url_joined_resultssnowflake_stage_joined_resultssnowflake_ocsp_joined_results"snowflake_telemetry_joined_resultsr   r   r   generate_reportw  sR   
z$ConnectionDiagnostic.generate_reportregistry_keyr   dict[str, str]c                 C  sN   i }d}	 zt ||}W n
 ty   Y |S w |dd ||d < |d }q)z%Gets values from windows registry keyr   Tr	   N)winreg	EnumValueOSError)rg   r,  registry_key_valuesr   registry_key_valuer   r   r   __get_win_registry_values  s   z.ConnectionDiagnostic.__get_win_registry_valuesr  registry_key_strc              	   C  s  |dkrt j}n
|dkrt j}nd}t ||}d}|durz	 t ||}|d }|rtj||}dt|	 v sCdt|	 v rt ||dt j
}	| |	}
d	|
 v r|
d	 d }| |d
| d| d|  d| d}z| jj|dd}|j}| |d|  W nT ty   Y nLw d|
 v r|
d d }| |d| d| d|  n/d|
 v r|
d d }|dkr| |d| d| d|  n| |d| d| d|
  | ||| q  ty   Y dS w dS )z@Walks the windows registry to search for key relating to proxiesr   r   Nr   Tr	   zinternet settingswpadAutoConfigURLziThere may be a proxy: Found a Wpad in Windows registry: Check proxy config for auto detect script: hkey: z : z	 : wpad: r   z	/wpad.datrt   )rv   z&Wpad request returned possible proxy: ProxyServerzFThere may be a proxy: Found a proxy server in Windows registry: hkey: z : ProxyServer: ProxyEnablez?There may be a proxy: Proxy is enabled per the registry: hkey: z : ProxyEnable: zCFound Wpad key in registry: Most likely nothing, but review: hkey: z	: value: )r.  r   r   OpenKeyEnumKeyrR   pathr
  r.   lowerKEY_READ._ConnectionDiagnostic__get_win_registry_valuesr  rJ   rf   r	  r  r   r  r0  )rg   rh   r  r4  hkeyregistryr   r,  new_registry_key_strnew_registry_keyvaluesr5  r   r  
proxy_infoproxyproxy_enabler   r   r   __walk_win_registry  s   

@z(ConnectionDiagnostic.__walk_win_registry)NNNNNNN)r-   r.   r/   r.   r0   r1   r2   r1   r3   r1   r4   r1   r5   r1   r6   r1   r7   r8   r9   r:   )rj   r.   r9   rl   )rs   rt   r=   )
r/   r.   rn   ru   rv   ru   rh   r.   r9   r.   )r9   r:   )r9   r   )r=   )r/   r.   rh   r.   r9   r:   )rs   r=   )r/   r.   rn   ru   rh   r.   r9   r:   )r   r   r9   r.   )rh   r.   r   r.   r9   r:   )r,  r   r9   r-  )rh   r.   r  r.   r4  r.   r9   r:   )__name__
__module____qualname____doc__rk   rT   r   r   r`   r   rU   r  rJ   r   r"  r+  r>  r  r   r   r   r   r,   9   s8    	 

^
0
g


~

Er,   )r   r   )r    r!   )(
__future__r   r   r   r(   rR   r   r   r   rX   r   loggingr   pathlibr   typingr   r   urllib.requestr   r   r   compatr
   r   rc   r   r7   r   r   url_utilr   vendoredr   rH  rL   r.  r   r+   r,   r   r   r   r   <module>   s6    

