
    ]j6<                        d 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m	Z	m
Z
 ddlmZ ddlmZ e
rddlZddlZ ej"                  e      Z ej(                         d        Z ej,                  d      Z ej,                  d	      Z ej,                  d
      Z ej,                  d      ZdZ G d de      Z G d de      Z G d dej>                        Z  G d de      Z!y)z2Utilities for Regional Access Boundary management.    N)
NamedTupleOptionalTYPE_CHECKING)_helpers)environment_varsc                      t         j                  j                  t        j                        } | y| j                         dv S )an  Checks if Regional Access Boundary is enabled via environment variable.

    The environment variable is interpreted as a boolean with the following
    (case-insensitive) rules:
    - "true", "1" are considered true.
    - Any other value (or unset) is considered false.

    Returns:
        bool: True if Regional Access Boundary is enabled, False otherwise.
    F)true1)osenvirongetr   "GOOGLE_AUTH_TRUST_BOUNDARY_ENABLEDlower)values    U/root/env/lib/python3.12/site-packages/google/auth/_regional_access_boundary_utils.py#is_regional_access_boundary_enabledr   #   s6     JJNN+NNOE};;=M))       )hours      )minuteszx-allowed-locationsc                       e Zd ZU dZee   ed<   eej                     ed<   eej                     ed<   ej                  ed<   y)_RegionalAccessBoundaryDataa  Data container for a Regional Access Boundary snapshot.

    Attributes:
        encoded_locations (Optional[str]): The encoded Regional Access Boundary string.
        expiry (Optional[datetime.datetime]): The hard expiration time of the boundary data.
        cooldown_expiry (Optional[datetime.datetime]): The time until which further lookups are skipped.
        cooldown_duration (datetime.timedelta): The current duration for the exponential cooldown.
    encoded_locationsexpirycooldown_expirycooldown_durationN)	__name__
__module____qualname____doc__r   str__annotations__datetime	timedelta r   r   r   r   H   sA      }$X&&''h//00)))r   r   c                   N    e Zd ZdZd Zd Zd Zd Zd ZddZ	d	 Z
d
 Zd Zd Zy)_RegionalAccessBoundaryManagerzManages the Regional Access Boundary state and its background refresh.

    The actual data is held in an immutable `_RegionalAccessBoundaryData` object
    and is swapped atomically to ensure thread-safe, lock-free reads.
    c                     t        d d d t              | _        t               | _        t        j                         | _        d| _        y )Nr   r   r   r   F)	r   )DEFAULT_REGIONAL_ACCESS_BOUNDARY_COOLDOWN_data%_RegionalAccessBoundaryRefreshManagerrefresh_manager	threadingLock_update_lock-_use_blocking_regional_access_boundary_lookupselfs    r   __init__z'_RegionalAccessBoundaryManager.__init___   s?    0" G	

  EF%NN,=B:r   c                 D    | j                   j                         }d|d<   |S )z9Pickle helper that serializes the _update_lock attribute.Nr2   __dict__copyr5   states     r   __getstate__z+_RegionalAccessBoundaryManager.__getstate__j   s#    ""$ $nr   c                 l    | j                   j                  |       t        j                         | _        y)z;Pickle helper that deserializes the _update_lock attribute.N)r9   updater0   r1   r2   r;   s     r   __setstate__z+_RegionalAccessBoundaryManager.__setstate__p   s#    U#%NN,r   c                     t        |t              st        S | j                  |j                  k(  xr | j                  |j                  k(  S )z!Checks if two managers are equal.)
isinstancer)   NotImplementedr-   r3   )r5   others     r   __eq__z%_RegionalAccessBoundaryManager.__eq__u   sF    %!?@!!JJ%++% CBBBBC	
r   c                     d| _         y)zEnables blocking Regional Access Boundary lookup.

        When enabled, the Regional Access Boundary lookup will be performed
        synchronously in the calling thread instead of asynchronously in a
        background thread.
        TN)r3   r4   s    r   enable_blocking_lookupz5_RegionalAccessBoundaryManager.enable_blocking_lookup   s     >B:r   Nc                 <    |sd}t        ||dt              | _        y)aI  Manually sets the regional access boundary to the client provided seed.

        Args:
            encoded_locations (Optional[str]): The encoded locations string.
            expiry (Optional[datetime.datetime]): The expiry time for the boundary.
                If encoded_locations is not provided, expiry is ignored.
        Nr+   )r   r,   r-   )r5   r   r   s      r   $set_initial_regional_access_boundaryzC_RegionalAccessBoundaryManager.set_initial_regional_access_boundary   s$     !F0/ G	

r   c                     | j                   }|j                  rA|j                  5t        j                         |j                  k  r|j                  |t
        <   y|j                  t
        d       y)a\  Applies the Regional Access Boundary header to the provided dictionary.

        If the boundary is valid, the 'x-allowed-locations' header is added
        or updated. Otherwise, the header is removed to ensure no stale
        data is sent.

        Args:
            headers (MutableMapping[str, str]): The headers dictionary to update.
        N)r-   r   r   r   utcnow _REGIONAL_ACCESS_BOUNDARY_HEADERpop)r5   headersrab_datas      r   apply_headersz,_RegionalAccessBoundaryManager.apply_headers   sR     ::%%OO'HOO,=,O8@8R8RG45KK8$?r   c                 r   | j                   }|j                  r5|j                  r)t        j                         |j                  t
        z
  k  ry|j                  r"t        j                         |j                  k  ry| j                  r| j                  ||       y| j                  j                  |||        y)a  Starts a background thread to refresh the Regional Access Boundary if needed.

        Args:
            credentials (google.auth.credentials.Credentials): The credentials to refresh.
            request (google.auth.transport.Request): The object used to make HTTP requests.
        N)r-   r   r   r   rK   *REGIONAL_ACCESS_BOUNDARY_REFRESH_THRESHOLDr   r3   start_blocking_refreshr/   start_refresh)r5   credentialsrequestrO   s       r   maybe_start_refreshz2_RegionalAccessBoundaryManager.maybe_start_refresh   s     :: &&!!KKM  ##(9H<T<T(T ==''W=  ..{GTJr   c                     	 |j                  |d      }| j                  |       y# t        $ r=}t        j                  t              rt        j                  d|d       d}Y d}~Sd}~ww xY w)a  Initiates a blocking lookup of the Regional Access Boundary.

        If the lookup raises an exception, it is caught and logged as a warning,
        and the lookup is treated as a failure (entering cooldown). Exceptions
        are not propagated to the caller.

        Args:
            credentials (google.auth.credentials.Credentials): The credentials to refresh.
            request (google.auth.transport.Request): The object used to make HTTP requests.
        T)	fail_fastz@Blocking Regional Access Boundary lookup raised an exception: %sexc_infoN) _lookup_regional_access_boundary	Exceptionr   is_logging_enabled_LOGGERwarning%process_regional_access_boundary_info)r5   rU   rV   regional_access_boundary_infoes        r   rS   z5_RegionalAccessBoundaryManager.start_blocking_refresh   su    	1
 <<WPT<U * 	223PQ  	1**73V!   
 -1)	1s   ' 	A-3A((A-c                    | j                   5  | j                  }|rl|j                  d      }t        |t	        j
                         t        z   dt              }t	        j                  t              rt        j                  d       nt	        j                  t              rt        j                  d       t	        j
                         |j                  z   }t        |j                  dz  t              }|j                  r&t	        j
                         |j                  kD  rd}d}n|j                   }|j                  }t        ||||      }|| _        ddd       y# 1 sw Y   yxY w)zProcesses the regional access boundary info and updates the state.

        Args:
            regional_access_boundary_info (Optional[Mapping[str, str]]): The regional access
                boundary info to process.
        encodedLocationsNr+   z+Regional Access Boundary lookup successful.z:Regional Access Boundary lookup failed. Entering cooldown.   )r2   r-   r   r   r   rK   $DEFAULT_REGIONAL_ACCESS_BOUNDARY_TTLr,   r^   r_   debugr`   r   min%MAX_REGIONAL_ACCESS_BOUNDARY_COOLDOWNr   r   )	r5   rb   current_datar   updated_datanext_cooldown_expirynext_cooldown_durationnext_encoded_locationsnext_expirys	            r   ra   zD_RegionalAccessBoundaryManager.process_regional_access_boundary_info   s>     1	&::L,$A$E$E&%!  ;&7#??,/SS$(&O	  ..w7MM"OP ..w7OOT
 OO%(F(FF % *- 22Q69*&  &&8??+<|?R?R+R-1*"&K-9-K-K*"."5"5K:&<&$8&<	  &DJc1	& 1	& 1	&s   EEE&)NN)r   r    r!   r"   r6   r=   r@   rE   rG   rI   rP   rW   rS   ra   r'   r   r   r)   r)   X   s>    	C-

B
$@&K8R:8&r   r)   c                   6     e Zd ZdZ	 	 	 	 	 	 d fdZd Z xZS )$_RegionalAccessBoundaryRefreshThreadzAThread for background refreshing of the Regional Access Boundary.c                 Z    t         |           d| _        || _        || _        || _        y )NT)superr6   daemon_credentials_request_rab_manager)r5   rU   rV   rab_manager	__class__s       r   r6   z-_RegionalAccessBoundaryRefreshThread.__init__$  s.     	''r   c                    	 | j                   j                  | j                        }| j                  j                  |       y# t        $ r=}t	        j
                  t              rt        j                  d|d       d}Y d}~]d}~ww xY w)a  
        Performs the Regional Access Boundary lookup and updates the state.

        This method is run in a separate thread. It delegates the actual lookup
        to the credentials object's `_lookup_regional_access_boundary` method.
        Based on the lookup's outcome (success or complete failure after retries),
        it updates the cached Regional Access Boundary information,
        its expiry, its cooldown expiry, and its exponential cooldown duration.
        zDAsynchronous Regional Access Boundary lookup raised an exception: %sTrZ   N)
rv   r\   rw   r]   r   r^   r_   r`   rx   ra   )r5   rb   rc   s      r   runz(_RegionalAccessBoundaryRefreshThread.run0  s    	1!!BB4==Q * 	??)	
  	1**73Z!   
 -1)	1s   %A 	B	3BB	)rU   z=google.auth.credentials.CredentialsWithRegionalAccessBoundaryrV   zgoogle.auth.transport.Requestry   r)   )r   r    r!   r"   r6   r|   __classcell__)rz   s   @r   rr   rr   !  s,    K
(T
( 1
( 6	
(
r   rr   c                   (    e Zd ZdZd Zd Zd Zd Zy)r.   zKManages a thread for background refreshing of the Regional Access Boundary.c                 D    t        j                         | _        d | _        y )N)r0   r1   _lock_workerr4   s    r   r6   z._RegionalAccessBoundaryRefreshManager.__init__R  s    ^^%
r   c                 N    | j                   j                         }d|d<   d|d<   |S )z?Pickle helper that serializes the _lock and _worker attributes.Nr   r   r8   r;   s     r   r=   z2_RegionalAccessBoundaryRefreshManager.__getstate__V  s,    ""$gir   c                 z    | j                   j                  |       t        j                         | _        d| _        y)zAPickle helper that deserializes the _lock and _worker attributes.N)r9   r?   r0   r1   r   r   r;   s     r   r@   z2_RegionalAccessBoundaryRefreshManager.__setstate__]  s)    U#^^%
r   c                    | j                   5  | j                  r$| j                  j                         r
	 ddd       y	 t        j                  |      }t        |||      | _        | j                  j                          ddd       y# t
        $ rA}t        j                  t              rt        j                  d|       Y d}~ddd       yd}~ww xY w# 1 sw Y   yxY w)a  
        Starts a background thread to refresh the Regional Access Boundary if one is not already running.

        Args:
            credentials (CredentialsWithRegionalAccessBoundary): The credentials
                to refresh.
            request (google.auth.transport.Request): The object used to make
                HTTP requests.
            rab_manager (_RegionalAccessBoundaryManager): The manager container to update.
        NzCould not deepcopy transport for background RAB refresh. Skipping background refresh to avoid thread safety issues. Exception: %s)r   r   is_aliver:   deepcopyr]   r   r^   r_   r`   rr   start)r5   rU   rV   ry   copied_requestrc   s         r   rT   z3_RegionalAccessBoundaryRefreshManager.start_refreshc  s     ZZ 	!|| 5 5 7	! 	!

!%w!7 @^[DL LL )	! 	!  ..w7OO( 	 	! 	!	! 	!s4   (CB	,C		C/CCCCCN)r   r    r!   r"   r6   r=   r@   rT   r'   r   r   r.   r.   O  s    U!r   r.   )"r"   r:   r%   	functoolsloggingr   r0   typingr   r   r   google.authr   r   google.auth.credentialsgooglegoogle.auth.transport	getLoggerr   r_   	lru_cacher   r&   rg   rR   r,   rj   rL   r   objectr)   Threadrr   r.   r'   r   r   <module>r      s    9     	  6 6   (" 
'

H
% * *& (:x'9'9'B $ .@X-?-?a-H * -?H,>,>r,J ) );(:(:(C % $9  ** * F&V F&R+
9+;+; +
\3!F 3!r   