
    l
j)/                         d dl Z d dlZd dlZ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 G d d      Z G d d      Zd Zd Zd	Z G d
 d      Z G d d      Zd Zd Z G d d      Zy)    N)futuresc                       e Zd Zd Zd Zy)
RegRequestc                     || _         || _        || _        || _        || _        d | _        || _        d | _        d| _        y )NF)	api_versionidenthostnameip_infoportrequest	auth_portlocked_cert	cancelled)selfr   r	   r
   r   r   r   s          ./usr/libexec/warpinator/remote_registration.py__init__zRegRequest.__init__   sG    &
 	  #    c                     d| _         y NT)r   r   s    r   cancelzRegRequest.cancel"   s	    r   N)__name__
__module____qualname__r   r    r   r   r   r      s     r   r   c                   $    e Zd Zd Zd Zd Zd Zy)	Registrarc                     d | _         d | _        i | _        t        j                         | _        || _        || _        || _        | j                          y N)
reg_server_v1reg_server_v2active_registrations	threadingLockreg_lockr
   r   r   start_registration_servers)r   r
   r   r   s       r   r   zRegistrar.__init__&   sI    !!$&!!(	"'')r   c                    | j                   | j                   j                          | j                  1| j                  j                  d      j                          d | _        t	        j
                  d| j                  | j                  fz         t        | j                  | j                        | _         t	        j
                  d| j                  | j                  fz         t        | j                  | j                        | _        y )N   gracez1Starting v1 registration server (%s) with port %dz6Starting v2 registration server (%s) with auth port %d)r    stopr!   waitloggingdebugr
   r   RegistrationServer_v1r   RegistrationServer_v2r   s    r   r&   z$Registrar.start_registration_servers2   s    )##%)##!#,113!%DIT\\[_[d[dLeef24<<KNRVR^R^`d`n`nQoop24<<Pr   c                    | j                   5  | j                  j                         D ]  }| j                  |   j                          ! i | _        d d d        | j                  6t        j                  d       | j                  j                          d | _        | j                  r7t        j                  d       | j                  j                          d | _        y y # 1 sw Y   xY w)Nz Stopping v1 registration server.z Stopping v2 registration server.)	r%   r"   keysr   r    r-   r.   r+   r!   )r   keys     r   shutdown_registration_serversz'Registrar.shutdown_registration_servers?   s    ]] 	+00557 8))#.5578(*D%	+
 )MM<=##%!%DMM<=##%!%D 	+ 	+s   ACC(c                 J   t        ||||||      }| j                  5  || j                  |<   d d d        d }|dk(  rt        |      }n|dk(  rt	        |      }| j                  5  	 | j                  |= d d d        |S # 1 sw Y   QxY w# t
        $ r Y !w xY w# 1 sw Y   |S xY w)N12)r   r%   r"   register_v1register_v2KeyError)	r   r   r	   r
   r   r   r   detailsrets	            r   registerzRegistrar.registerO   s    UHgtYT]] 	7/6D%%e,	7 #g&CCg&C]] 	--e4	 
#	7 	7  		 
s5   A=$B&B	=B		BBBBB"N)r   r   r   r   r&   r4   r=   r   r   r   r   r   %   s    
*Q& r   r   c                 :   t        j                  d| j                  | j                  | j                  fz         t        |       }|t        j                  j                  k(  r:t        j                  d| j                  | j                  | j                  fz         yy)Nz+Registering with %s (%s:%d) - api version 1z2Unable to register with %s (%s:%d) - api version 1FT)	r-   r.   r	   r
   r   retrieve_remote_certutilCertProcessingResultFAILURE)r;   successs     r   r8   r8   f   s    
 MM?7CSCSU\UdUdfmfrfrBsst"7+G$++333J ' 0 0'//7<<PQ 	Rr   c                    t        j                  d| j                  | j                  | j                  fz         t        | j                  | j                        | _        | j                  j                         }|| j                  rt        j                  j                  S t        j                         j                  | j                  | j                  |      S )Nz1Auth: Starting a new RequestLoop for '%s' (%s:%d))r-   r.   r	   r
   r   Requestr   r   r@   rA   rB   authget_singletonprocess_remote_cert)r;   datas     r   r?   r?   v   s    MMEIYIY[b[j[jlslxlxHyyzgoow||<GO??""$D|w((((00033G4D4D4;OO48: :r   s   REQUESTc                       e Zd Zd Zd Zy)rE   c                      || _         || _        y r   )r
   r   r   r
   r   s      r   r   zRequest.__init__   s    	r   c                    t        j                  d| j                  | j                  fz         | j                  j	                         \  }}}	 |t
        j                  k(  r|nd|d}t        j
                  |t
        j                        }|j                  d       |j                  t        || j                  f       |j                  d      \  }}||| j                  fk(  r|S 	 y # t
        j                  $ r1 t        j                  d| j                  | j                  fz         Y y t
        j                  $ r/}t        j                  d|| j                  |fz         Y d }~y d }~ww xY w)Nz)Auth: Requesting cert from remote (%s:%d)[]g      @  zKAuth: Cert request failed from remote (%s:%d) - (Is their udp port blocked?z+Something wrong with cert request (%s:%s): )r-   r.   r
   r   get_usable_ipsocketAF_INET
SOCK_DGRAM
settimeoutsendtoREQUESTrecvfromtimeouterrorcritical)	r   	remote_ip_
ip_versionipserver_sockreplyaddres	            r   r   zRequest.request   sC   AT\\SWS\S\D]]^#'<<#=#=#? 	1j	h(FNN::)@UB --
F4E4EFK""3'wTYY8%..t4KE4	499-- .  ~~ 	;MMg!%tyy 9: ;
  || 	hJiY]YbYbdeMffgg	hs    BC) )AE-,E->%E((E-N)r   r   r   r   r   r   r   r   rE   rE      s    r   rE   c                       e Zd Zd Zd Zd Zy)r/   c                 j   d| _         || _        || _        t        j                  | j
                  t        j                  f      | _        t        j                  | j
                  t        j                  f      | _
        | j                  j                          | j                  j                          y )NF)targetargs)exitr
   r   r#   Threadserve_cert_threadrR   rS   thread4AF_INET6thread6startrL   s      r   r   zRegistrationServer_v1.__init__   sy    		 ''t/E/EV^^L]^ ''t/E/EV__L^_r   c                    d }|t         j                  k(  r| j                  j                  }n)|t         j                  k(  r| j                  j
                  }|	 t        j                   |t         j                        }|j                  d       |j                  || j                  f       	 	 |j                  d      \  }}|t        k(  r4t!        j"                         j%                         }|j'                  ||       Sy # t         j                  $ r+}t        j                  dt        |      z         Y d }~y d }~ww xY w# t         j(                  $ r+}| j*                  r|j-                          Y d }~y Y d }~|d }~ww xY w)Ng      ?z1Could not create udp socket for cert requests: %srP   )rR   rS   r
   ip4_addressrl   ip6_addressrT   rU   bindr   rZ   r-   r[   strrX   rW   rF   rG   get_encoded_local_certrV   rY   rh   close)r   r^   local_ipr`   rc   rI   address	cert_datas           r   rj   z'RegistrationServer_v1.serve_cert_thread   s=   '||//H6??*||//H$mmJ8I8IJ&&s+  (DII!67
 	$/$8$8$>MD'w$($6$6$8$O$O$Q	#**9g>   
 <<   !TWZ[\W]!]^ ~~ yy#))+ !s2   AD  -AE  D>!D99D>E?E::E?c                 z    d| _         | j                  j                          | j                  j                          y r   )rh   rk   joinrm   r   s    r   r+   zRegistrationServer_v1.stop   s)    	r   N)r   r   r   r   rj   r+   r   r   r   r/   r/      s    8r   r/   c                    t        j                  d| j                  | j                  | j                  fz         d }t        j                  t        | fdt        z        }t        j                  d| j                  | j                  | j                  | j                  fz         |j                          |j                          | j                  O| j                  sCt        j                         j!                  | j                  | j                  | j                        }|t"        j$                  j&                  k(  r;t        j                  d| j                  | j                  | j                  fz         |S |t"        j$                  j(                  k(  r;t        j                  d| j                  | j                  | j                  fz         |S |t"        j$                  j*                  k(  r;t        j                  d| j                  | j                  | j                  fz         |S |t"        j$                  j,                  k(  r9t        j                  d| j                  | j                  | j                  fz         |S )	Nz+Registering with %s (%s:%d) - api version 2zremote-auth-thread-%s)rf   rg   namez&remote-registration-thread-%s-%s:%d-%sz2Unable to register with %s (%s:%d) - api version 2z7Successfully registered with %s (%s:%d) - api version 2zASuccessfully updated registration with %s (%s:%d) - api version 2zLCertificate already up to date, nothing to do for %s (%s:%d) - api version 2)r-   r.   r	   r
   r   r#   ri   register_with_remote_threadidr   rn   rz   r   r   rF   rG   rH   r@   rA   rB   CERT_INSERTEDCERT_UPDATEDCERT_UP_TO_DATE)r;   rC   remote_threads      r   r9   r9      s+    MM?7CSCSU\UdUdfmfwfwBxxyG$$,Gwj_vy{_{|MMM:g>N>NPWP_P_aharart{  uB  uB  >C  C  D&w/@/@$$&::7;K;K;B??;B;N;NP $++333J ' 0 0'//7CTCTUV 	W N 
D--;;	;O ' 0 0'//7CTCTUV 	W N 
D--::	:Y ' 0 0'//7CTCTUV 	W
 N 
D--==	=d ' 0 0'//7CTCTUV 	WNr   c           	         t        j                  d| j                  d| j                  d       | j                  j	                         \  }}}|t
        j                  k(  r|nd|d}t        j                  d|| j                  fz        5 }t        j                  |      }	 t        j                  |      }|j                  t        j                  |t!        j"                               d	      }|j$                  j'                  d
      | _        d d d        y # t(        $ rT}|j+                          t        j,                  d| j                  | j                  | j                  |fz         Y d }~ad }~ww xY w# 1 sw Y   y xY w)NzRemote: Attempting to register z ()rN   rO   %s:%d)r_   r	      )rY   zutf-8zGProblem with remote registration thread: %s (%s:%d) - api version 2: %s)r-   r.   r	   r
   rQ   rR   rS   grpcinsecure_channelr   channel_ready_futurewarp_pb2_grpcWarpRegistrationStubRequestCertificatewarp_pb2r   r@   get_hostnamer   encode	Exceptionr   r[   )	r;   r\   rv   r^   channelfuturestubr<   rc   s	            r   r}   r}      sU   MMg>N>NPWP_P_`a&-oo&C&C&E#Ix'6>>9	9?TI			w)W5F5F)GG	H RG**73
	R 55g>D))(*=*=UYUfUfUh*i23 * 5C"%//"8"8"AGR R  	RMMOf(('//7;L;LaPQ R R	RR Rs2   E=%A/D	E:&A
E50E=5E::E==Fc                   D    e Zd Zd Zd Zd Zd Zdej                  fdZ	y)r0   c                     d| _         || _        || _        d | _        t	        j
                         | _        t	        j                  | j                        | _	        | j                  j                          y )NF)rf   )rh   r
   r   serverr#   Eventserver_thread_keepaliveri   rj   threadrn   )r   r
   r   s      r   r   zRegistrationServer_v2.__init__	  sU    	"'0'8$&&d.D.DEr   c                 l   t        j                  t        j                  d            | _        t	        j
                  | | j                         | j                  j                  >| j                  j                  d| j                  j                  | j                  fz         | j                  j                  >| j                  j                  d| j                  j                  | j                  fz         | j                  j                          | j                  j                         s6| j                  j                  d       | j                  j                         s6t        j                   d       | j                  j#                  d      j                          t        j                   d       y )	Nr(   )max_workersr   z[%s]:%d
   zRegistration Server v2 stoppingr)   zRegistration Server v2 stopped)r   r   r   ThreadPoolExecutorr   &add_WarpRegistrationServicer_to_serverr
   rp   add_insecure_portr   rq   rn   r   is_setr,   r-   r.   r+   r   s    r   rj   z'RegistrationServer_v2.serve_cert_thread  s,   kk'"<"<"KL<<T4;;O<<##/KK))'T\\5M5Mt~~4^*^_<<##/KK)))t||7O7OQUQ_Q_6`*`a..557((--b1 ..557 	78q!&&(67r   c                 l    | j                   j                          | j                  j                          y r   )r   setr   rz   r   s    r   r+   zRegistrationServer_v2.stop%  s$    $$((*r   c                     t        j                  d|j                  d|j                  d       t	        j
                  t        j                         j                               S )Nz1Registration Server RPC: RequestCertificate from z '')r   )	r-   r.   r	   r_   r   RegResponserF   rG   rt   )r   r   contexts      r   r   z(RegistrationServer_v2.RequestCertificate)  sD    T[TdTdfmfpfpqr##0B0B0D0[0[0]^^r   regc           	         t        j                  d|j                  z          | j                  |       t	        j
                  t        j                         | j                  j                  t        j                         t        j                         t        t        j                        | j                   | j                  j"                        S )Nz"Received manual registration from )
service_idr_   r   r	   r   r   ipv6)r-   r.   r   service_registration_handlerr   ServiceRegistrationprefsget_connect_idr
   rp   get_portr@   r   intconfigRPC_API_VERSIONr   rq   )r   r   r   s      r   RegisterServicez%RegistrationServer_v2.RegisterService.  s    :S^^KL))#.++u7K7K7M/3||/G/G161A595F5F5H8;F<R<R8S6:nn151I1IK 	Kr   N)
r   r   r   r   rj   r+   r   r   r   r   r   r   r   r0   r0     s+    	8"_
	K(">"> 	Kr   r0   )r-   r#   rR   
concurrentr   r   r   r   rF   r@   r   r   r   r   r8   r?   rW   rE   r/   r9   r}   r0   r   r   r   <module>r      s{               (= =B :  :* *`@R*/K /Kr   