
    Oǻid0                         d dl 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	m
Z
mZ d dlmZmZ d dlmZ d dlmZ  ej&                  e      Z G d d	      Z G d
 d      Z G d d      Z G d d      Zd Zy)    N)datetimetimezone)sleep)Any	AwaitableCallableUnion)RequestTokenErrTokenRenewalErr)IdentityProviderInterface)TokenResponsec                       e Zd ZdZd Zedeeegdf   e	f   fd       Z
e
j                  deeegdf   e	f   ddfd       Z
edeeegdf   e	f   fd       Zej                  deeegdf   e	f   ddfd	       Zy)
CredentialsListenerzz
    Listeners that will be notified on events related to credentials.
    Accepts callbacks and awaitable callbacks.
    c                      d | _         d | _        y N)_on_next	_on_errorselfs    V/opt/lhia/marcimex/agent/venv/lib/python3.12/site-packages/redis/auth/token_manager.py__init__zCredentialsListener.__init__   s        returnNc                     | j                   S r   r   r   s    r   on_nextzCredentialsListener.on_next   s    }}r   callbackc                     || _         y r   r   r   r   s     r   r   zCredentialsListener.on_next   s	     r   c                     | j                   S r   r   r   s    r   on_errorzCredentialsListener.on_error!   s    ~~r   c                     || _         y r   r!   r   s     r   r"   zCredentialsListener.on_error%   s	    !r   )__name__
__module____qualname____doc__r   propertyr	   r   r   r   r   setter	Exceptionr"    r   r   r   r      s    
 xt4i?@   ^^!hud{&;Y&F G !D ! ! %)d): ;Y FG   __"xT0A'BI'M!N "SW " "r   r   c                   4    e Zd ZdedefdZdefdZdefdZy)RetryPolicymax_attemptsdelay_in_msc                      || _         || _        y r   )r.   r/   )r   r.   r/   s      r   r   zRetryPolicy.__init__+   s    (&r   r   c                     | j                   S )zW
        Retry attempts before exception will be thrown.

        :return: int
        )r.   r   s    r   get_max_attemptszRetryPolicy.get_max_attempts/   s        r   c                     | j                   S )zI
        Delay between retries in seconds.

        :return: int
        )r/   r   s    r   get_delay_in_mszRetryPolicy.get_delay_in_ms7   s     r   N)r$   r%   r&   intfloatr   r2   r4   r+   r   r   r-   r-   *   s-    'S 'u '!# !   r   r-   c                   T    e Zd ZdedededefdZdefdZdefdZdefd	Z	defd
Z
y)TokenManagerConfigexpiration_refresh_ratiolower_refresh_bound_millis%token_request_execution_timeout_in_msretry_policyc                 <    || _         || _        || _        || _        y r   )_expiration_refresh_ratio_lower_refresh_bound_millis&_token_request_execution_timeout_in_ms_retry_policy)r   r9   r:   r;   r<   s        r   r   zTokenManagerConfig.__init__A   s*     *B&+E(1 	3 *r   r   c                     | j                   S )a&  
        Represents the ratio of a token's lifetime at which a refresh should be triggered. # noqa: E501
        For example, a value of 0.75 means the token should be refreshed
        when 75% of its lifetime has elapsed (or when 25% of its lifetime remains).

        :return: float
        )r>   r   s    r   get_expiration_refresh_ratioz/TokenManagerConfig.get_expiration_refresh_ratioO   s     ---r   c                     | j                   S )a  
        Represents the minimum time in milliseconds before token expiration
        to trigger a refresh, in milliseconds.
        This value sets a fixed lower bound for when a token refresh should occur,
        regardless of the token's total lifetime.
        If set to 0 there will be no lower bound and the refresh will be triggered
        based on the expirationRefreshRatio only.

        :return: int
        )r?   r   s    r   get_lower_refresh_bound_millisz1TokenManagerConfig.get_lower_refresh_bound_millisZ   s     ///r   c                     | j                   S )z
        Represents the maximum time in milliseconds to wait
        for a token request to complete.

        :return: int
        )r@   r   s    r   )get_token_request_execution_timeout_in_msz<TokenManagerConfig.get_token_request_execution_timeout_in_msg   s     :::r   c                     | j                   S )z_
        Represents the retry policy for token requests.

        :return: RetryPolicy
        )rA   r   s    r   get_retry_policyz#TokenManagerConfig.get_retry_policyp   s     !!!r   N)r$   r%   r&   r6   r5   r-   r   rC   rE   rG   rI   r+   r   r   r8   r8   @   s]    *"'* %(* 03	*
 "*	.e 	.0 0;3 ;"+ "r   r8   c                       e Zd ZdedefdZd Z	 ddedede	g df   fd	Z
	 	 	 dded
ededede	g df   f
dZd ZddefdZddefdZdededefdZdefdZdedefdZddefdZ	 ddedej,                  fdZy)TokenManageridentity_providerconfigc                 X    || _         || _        d | _        d | _        d | _        d| _        y )Nr   )_idp_config_next_timer	_listener_init_timer_retries)r   rL   rM   s      r   r   zTokenManager.__init__z   s0     &	r   c                 N    t         j                  d       | j                          y )NzToken manager are disposed)loggerinfostopr   s    r   __del__zTokenManager.__del__   s    01		r   listenerskip_initialr   Nc                     | _         	 t        j                         t        j                          fd}j                  |       _        t        j                  d       j                           j                  S # t        $ rh t        j                         t        j                         fd}t        j                  |d      }|j                          j                          Y w xY w)Nc                  p    t        j                          j                           j                          y r   )asyncioset_event_loopsetrun_forever)loop
loop_readys   r   
start_loopz&TokenManager.start.<locals>.start_loop   s)    
 &&t,   "r   T)targetdaemonc                  t    	 j                          j                          y #  j                          w xY wr   )_renew_tokenr`   )	init_doner   r[   s   r   renew_with_callbackz/TokenManager.start.<locals>.renew_with_callback   s)     !!,/	s   % 7Token manager started)rR   r^   get_running_loopRuntimeErrornew_event_loop	threadingEventThreadstartwaitcall_soon_threadsaferS   rV   rW   rX   )	r   rZ   r[   rd   threadrj   ri   rb   rc   s	   ` `   @@@r   rr   zTokenManager.start   s    
 "	++-D0 OO%		   445HI+, 	yyK  	))+D #*J# %%ZEFLLN OO)	s   B A.C54C5block_for_initialinitial_delay_in_msc                 F  K   || _         t        j                         }t        j                         }t	        || j
                  ||      }|j                  |dz  |      | _        t        j                  d       |r|j                          d {    | j                  S 7 w)N  rk   )rR   r^   rl   rp   _async_to_sync_wrapper_renew_token_async
call_laterrS   rV   rW   rs   rX   )r   rZ   rv   rw   r[   rb   
init_eventwrappeds           r   start_asynczTokenManager.start_async   s      "'')]]_
 )$))<
  ??+>+EwO+,//###yy $s   BB!BB!c                     | j                   | j                   j                          | j                  | j                  j                          y y r   )rS   cancelrQ   r   s    r   rX   zTokenManager.stop   sC    '##%'##% (r   c                    	 | j                   j                  |      }d| _        t        |      S # t        $ r}| j                  | j                  j                         j                         k  r_| xj                  dz  c_        t        | j                  j                         j                         dz         | j                  |      cY d }~S |d }~ww xY wN   ry   r   )rO   request_tokenr
   rT   rP   rI   r2   r   r4   acquire_tokenr   r   force_refreshtokenes       r   r   zTokenManager.acquire_token   s    	II++M:E U##  	}}t||<<>OOQQ"dll335EEG$NO))-88	s!   / 	CBCCCCc                   K   	 | j                   j                  |      }d| _        t        |      S # t        $ r}| j                  | j                  j                         j                         k  r{| xj                  dz  c_        t        j                  | j                  j                         j                         dz         d {  7   | j                  |       d {  7  cY d }~S |d }~ww xY wwr   )rO   r   r
   rT   rP   rI   r2   r^   r   r4   acquire_token_asyncr   r   s       r   r   z TokenManager.acquire_token_async   s     
	II++M:E U##  	}}t||<<>OOQQ"mmLL113CCEL   "55mDDDD	sO   C41 C4	C1BC,CC,C" C,$C1%C4*C,,C11C4expire_date
issue_datec                 x    | j                  |      }| j                  ||      }t        ||      }|dk  rdS |dz  S )Nr   ry   )_delay_for_lower_refresh_delay_for_ratio_refreshmin)r   r   r   delay_for_lower_refreshdelay_for_ratio_refreshdelays         r   _calculate_renewal_delayz%TokenManager._calculate_renewal_delay   sI    "&"?"?"L"&"?"?Z"X+-DEAIq/54</r   c                     || j                   j                         z
  t        j                  t        j
                        j                         dz  z
  S Nry   )rP   rE   r   nowr   utc	timestamp)r   r   s     r   r   z%TokenManager._delay_for_lower_refresh   sB    ll99;<||HLL)335<>	
r   c                     ||z
  }||| j                   j                         z  z
  }||z
  t        j                  t        j
                        j                         dz  z
  S r   )rP   rC   r   r   r   r   r   )r   r   r   	token_ttlrefresh_befores        r   r   z%TokenManager._delay_for_ratio_refresh  s`    *,	"AACC

 ||HLL)335<>	
r   c                 4   	 | j                  d      }| j                  |j                         j                         |j                         j	                               }|j                         j                         rt        d      | j                  j                  t        j                  d       y|s*	 | j                  j                  |j                                |dk  ryt        j                         }|j                  || j                        | _        t        j#                  d| d       |S # t        $ r}t        |      d}~ww xY w# t        $ r=}| j                  j$                  || j                  j%                  |       Y d}~yd}~ww xY w)	zq
        Task to renew token from identity provider.
        Schedules renewal tasks based on token TTL.
        Tr   Requested token is expiredN@No registered callback for token renewal task. Renewal cancelledr    Next token renewal scheduled in  seconds)r   r   	get_tokenget_expires_at_msget_received_at_ms
is_expiredr   rR   r   rV   warningr*   r^   rl   r|   rh   rQ   rW   r"   )r   r[   	token_resr   r   rb   s         r   rh   zTokenManager._renew_token  se   !	'***>I11##%779##%88:E
 ""$//1%&BCC~~%%-V -NN**9+>+>+@A z++-D#ud6G6GHDKK:5'JK ! -)!,,-  	'~~&&.NN##A&&		'sI   B0E 3E 6)D5 E %AE 5	E>E		EE 	F3FFr}   c                 *  K   	 | j                  d       d{   }| j                  |j                         j                         |j                         j	                               }|j                         j                         rt        d      | j                  j                  *t        j                  d       	 |r|j                          yy|s2	 | j                  j                  |j                                d{    |dk  r	 |r|j                          yyt        j                         }t        || j                         }t        j#                  d| d       |j%                  ||       |r|j                          yy7 o7 # t        $ r}t        |      d}~ww xY w# t        $ rF}| j                  j&                  || j                  j'                  |       d{  7   Y d}~d}~ww xY w# |r|j                          w w xY ww)	zx
        Async task to renew tokens from identity provider.
        Schedules renewal tasks based on token TTL.
        Tr   Nr   r   r   r   r   )r   r   r   r   r   r   r   rR   r   rV   r   r`   r*   r^   rl   rz   r{   rW   r|   r"   )r   r[   r}   r   r   r   rb   r~   s           r   r{   zTokenManager._renew_token_async7  s    $	!"66T6JJI11##%779##%88:E
 ""$//1%&BCC~~%%-V *   '  -..001D1D1FGGG z    ++-D,T43J3JKGKK:5'JKOOE7+   E K" H  -)!,,-  	-~~&&...))!,,,		-   s   HF) FB"F) =HF) ,F  FF F) HAF) 4HF) F 	F&F!!F&&F) )	G826G3(G+)G3.G; 3G88G; ;HH)F)Fr   F)FN)r$   r%   r&   r   r8   r   rY   r   boolr   rr   r6   r   rX   r   r   r   r   r   r   rh   r^   rp   r{   r+   r   r   rK   rK   y   s   !:DV #.%. . 
"d(		.f #(%&"%   #	
  
"d(	0&$M $$ $ 0E 0u 0QV 0
E 


E 

u 

'' ''T GK,! ,!6=mm,!r   rK   c                       fd}|S )a  
    Wraps an asynchronous function so it can be used with loop.call_later.

    :param loop: The event loop in which the coroutine will be executed.
    :param coro_func: The coroutine function to wrap.
    :param args: Positional arguments to pass to the coroutine function.
    :param kwargs: Keyword arguments to pass to the coroutine function.
    :return: A regular function suitable for loop.call_later.
    c                  @    t        j                    i        y )N)rb   )r^   ensure_future)args	coro_funckwargsrb   s   r   r~   z'_async_to_sync_wrapper.<locals>.wrappedq  s    i88tDr   r+   )rb   r   r   r   r~   s   ```` r   rz   rz   f  s    E Nr   )r^   loggingro   r   r   timer   typingr   r   r   r	   redis.auth.errr
   r   redis.auth.idpr   redis.auth.tokenr   	getLoggerr$   rV   r   r-   r8   rK   rz   r+   r   r   <module>r      sf       '  2 2 ; 4 *			8	$" "6   ,6" 6"rj! j!Zr   