
    *                     D   S r SSKJr  SSKJr  SSKJr  SSKrSSKrSSKrSSKrSSK	r	SSK
Jr  SSK
Jr  SSKJr  SS	KJr  SS
KJr  SSKJr  SSKrSSKrSr " S S\R0                  5      r " S S\R0                  5      r " S S\R0                  5      r " S S\5      rg)z4WebSocket helper class for tunneling with Cloud IAP.    )absolute_import)division)unicode_literalsN) iap_tunnel_lightweight_websocket)iap_tunnel_websocket_utils)context_aware)
exceptions)log)encodingzbot:iap-tunnelerc                       \ rS rSrSrg)WebSocketConnectionClosed(    N__name__
__module____qualname____firstlineno____static_attributes__r       Alib/googlecloudsdk/api_lib/compute/iap_tunnel_websocket_helper.pyr   r   (       r   r   c                       \ rS rSrSrg)WebSocketInvalidOpcodeError,   r   Nr   r   r   r   r   r   ,   r   r   r   c                       \ rS rSrSrg)WebSocketSendError0   r   Nr   r   r   r   r   r   0   r   r   r   c                   n    \ rS rSrSr SS jrS rSS jrS rS r	S r
S	 rS
 rS rSS jrS rS rSrg)IapTunnelWebSocketHelper4   zEHelper class for common operations on websocket and related metadata.c	           	      x   XPl         X`l        X@l        Xl        S U l        [
        R                  " U5      n	[        R                  U	S.U l	        U(       a,  [        R                  U R                  S'   SU R                  S'   [        R                  " 5       n
U
(       a  U
R                  [        R                  R                  L a  [         R"                  " S5        OJU
R$                  n[         R"                  " SU5        XR                  S'   U
R&                  U R                  S'   [         R(                  " 5       [*        R,                  :w  a3  [*        R.                  " S	5      R1                  [*        R2                  5        SU l        S
U l        Xpl        U R8                  (       aL  [:        R<                  " UUU R>                  U R@                  U RB                  [
        RD                  /S9U l#        g [H        RJ                  " UUU R>                  U R@                  U RB                  [
        RD                  /S9U l#        g )N)	cert_reqsca_certsr#   Fcheck_hostnamez%Using enterprise certificate for mTLSzUsing client certificate %scertfilepassword	websocket )headeron_closeon_dataon_errorsubprotocols)&_on_data	_on_close_proxy_info_conn_id_receiving_threadutilsCheckCACertsFilesslCERT_REQUIRED_sslopt	CERT_NONEr   Configconfig_type
ConfigTypeENTERPRISE_CERTIFICATEr
   debugencrypted_client_cert_pathencrypted_client_cert_passwordGetVerbosityloggingDEBUG	getLoggersetLevelCRITICAL
_is_closed
_error_msg_should_use_new_websocketiap_websocketIapLightWeightWebsocket_OnClose_OnData_OnErrorSUBPROTOCOL_NAME
_websocketr(   WebSocketApp)selfurlheadersignore_certs
proxy_infor,   r+   should_use_new_websocketconn_idr$   
caa_config	cert_paths               r   __init__!IapTunnelWebSocketHelper.__init__7   s   MN!M!D%%l3H!$!2!2 (*DL"%--dll;',dll#$%%'J

 
 

"
"
9
9:		9:99			/;#,Z #-#L#LZ  W]]*$--g.>.>?DODO%="%%%==
==,,==../1do "..
==,,==../1dor   c                 $    U R                  5         g N)CloserR   s    r   __del__ IapTunnelWebSocketHelper.__del__j   s    JJLr   c                     U R                   (       d:   U R                  R                  5         U R                  (       d  Xl        SU l         gg!    N$= f)zClose the WebSocket.TN)rG   rP   closerH   )rR   msgs     r   r_   IapTunnelWebSocketHelper.Closem   sB    ?? __do s   A Ac                     U R                   =(       d2    U R                  =(       a    U R                  R                  5       (       + $ )z%Check to see if WebSocket has closed.)rG   r3   is_aliver`   s    r   IsClosed!IapTunnelWebSocketHelper.IsClosedx   s5    OO O##MD,B,B,K,K,M(MPr   c                     U R                   $ r^   )rH   r`   s    r   ErrorMsg!IapTunnelWebSocketHelper.ErrorMsg}   s    ??r   c           
          [         R                  " 5       [        R                  :X  a/  [         R                  " SU R
                  [        U5      USS 5        U R                  R                  U[        R                  R                  S9  g! [         a    U R                  5         e [        R                   a    U R                  5         [        5       e[          a  n[         R                  " SU R
                  SS9  ["        R$                  " 5       S   nU R                  5         [&        R(                  " [+        [,        R.                  " [1        U5      U5      US	95         SnAgSnAff = f)
z"Send data on WebSocket connection.z)[%d] SEND data_len [%d] send_data[:20] %rN   )opcodez1[%d] Error during WebSocket send of Data message.Texc_info   )tb)r
   rA   rB   rC   r>   r2   lenrP   sendr(   ABNFOPCODE_BINARYEnvironmentErrorr_   "WebSocketConnectionClosedExceptionr   	Exceptionsysrr   r	   reraiser   	tracebackformat_exception_onlytype)rR   	send_dataert   s       r   SendIapTunnelWebSocketHelper.Send   s   % 
			w}}	,		=--Y3B	B
oo9Y^^-I-IJ 
jjl77 (
jjl%'' 	%	iiC. <<>!b
jjl
Y<<T!WaH "$% %	%s   BB AE*BE%%E*c                    U R                   (       a  U R                  nOU R                  R                  nU(       a3  [        R                  " SU R
                  5         UR                  5         gg! [        [        R                  4 aP  n[        R                  " SU R
                  [        R                  " U5      5        U R                  5          SnAgSnAf  [        R                  " SU R
                  SS9  U R                  5          g= f)z)Send WebSocket Close message if possible.z
[%d] CLOSEz1[%d] Unable to send WebSocket Close message [%s].Nz2[%d] Error during WebSocket send of Close message.Trq   )rI   rP   sockr
   r>   r2   
send_closery   r(   rz   infosix	text_typer_   )rR   r   r   s      r   	SendClose"IapTunnelWebSocketHelper.SendClose   s     %%__d__!!d	iidmm,
  ::< Da 0	2

E	/

s   A0 0D
AC5Dc                     U R                   (       dO  [        R                  " U R                  S9U l        SU R                  l        U R                  R                  5         g g )N)targetT)rG   	threadingThread_ReceiveFromWebSocketr3   daemonstartr`   s    r   StartReceivingThread-IapTunnelWebSocketHelper.StartReceivingThread   sJ    ??(//++ -d&*d#
""$	 r   c                 T   Uc  Uc  U R                  5         gU< SU< 3n[        R                  " SU R                  U5        U R                  US9  US:X  a  g U R	                  5         g! [
        [        R                  4 a"    [        R                  " SU R                  SS9  e f = f)	z&Callback for WebSocket Close messages.Nz: z+[%d] Received WebSocket Close message [%s].)re   i  z)[%d] Error while processing Close messageTrq   )r_   r
   r   r2   r0   ry   r	   Error)rR   
close_codeclose_reason	close_msgs       r   rL   !IapTunnelWebSocketHelper._OnClose   s     l2 jjl&5IHH:]]I'JJ9JT 
nnj../ 	hh:}}t-s   A+ +<B'c           	      v   [         R                  " 5       [        R                  :X  a0  [         R                  " SU R
                  U[        U5      USS 5         U[        R                  R                  [        R                  R                  4;  a  [        SU-  5      eU R                  U5        g! [         aL  n[         R                  " SU R
                  [        R                   " U5      5        U R#                  5         e SnAf  [         R                  " SU R
                  SS9  U R#                  5         e = f)	z%Callback for WebSocket Data messages.z9[%d] RECV opcode [%r] data_len [%d] binary_data[:20] [%r]Nro   z!Unexpected WebSocket opcode [%r].z([%d] Error [%s] while sending to client.z)[%d] Error while processing Data message.Trq   )r
   rA   rB   rC   r>   r2   ru   r(   rw   OPCODE_CONTrx   r   r/   ry   r   r   r   r_   )rR   binary_datarp   unused_finishedr   s        r   rM    IapTunnelWebSocketHelper._OnData   s     W]]*	iiKvs;'7Sb9IK
 
	22!446 
6)*M*0+1 2 	2
mmK  	hh94==}}Q!
jjl	hh:DMM
jjls   AB1 1
D8;AD6D8c           
      @   U R                   (       d  [        R                  " SU R                  SS9  [        R                  " SSR                  [        R                  " [        U5      U5      5      -   U R                  5        [        R                  " U5      U l        g g )Nz'[%d] Error during WebSocket processing.Trq   z([%d] Error during WebSocket processing:
r)   )rG   r
   r>   r2   r   joinr~   r   r   r   r   rH   )rR   exception_objs     r   rN   !IapTunnelWebSocketHelper._OnError   s}     ??	ii9.	hh:wwy66tM7J7DF GG }} m4do r   c                     U R                   (       a  SnU R                   R                  (       d  U R                   R                  (       aT  [        R                  " U R                   R                  5      [        R                  " U R                   R                  5      4nU R
                  R                  [        U R                  U R                   R                  U R                   R                  US9  O)U R
                  R                  [        U R                  S9    U R                  5         g!    [        R                  " SU R                  SS9   N7!     N<= f= f!    [        R                  " SU R                  SS9   g!     g= f= f)z'Receive data from WebSocket connection.N)originsslopthttp_proxy_hosthttp_proxy_porthttp_proxy_auth)r   r   z*[%d] Error while receiving from WebSocket.Trq   z-[%d] Error while closing in receiving thread.)r1   
proxy_user
proxy_passr   DecoderP   run_foreverTUNNEL_CLOUDPROXY_ORIGINr8   
proxy_host
proxy_portr
   r   r2   r_   )rR   r   s     r   r   .IapTunnelWebSocketHelper._ReceiveFromWebSocket   s9   			&&$*:*:*E*E
 &__T-=-=-H-HI%__T-=-=-H-HIK/##+DLL ,,77 ,,77+	 	$ 	- 	##+C+/<< 	$ 	9
jjl	=	/
 	@	/sN   C0D. 3(D. E .E1 EEEEF! FFFF)
r2   rH   rG   r0   r/   r1   r3   rI   r8   rP   N)r   )r)   )r   r   r   r   __doc__r[   ra   r_   ri   rl   r   r   r   rL   rM   rN   r   r   r   r   r   r    r    4   sH    M 2311f	P
%4.%<8
5%r   r    )r   
__future__r   r   r   rB   r6   r|   r   r~   googlecloudsdk.api_lib.computer   rJ   r   r4   googlecloudsdk.corer   r	   r
   googlecloudsdk.core.utilr   r   r(   r   r   r   r   r   objectr    r   r   r   <module>r      s     ; &  '  
 
   \ N - * # - 
 - 
 0 0 *"2"2 )) ov or   