
    ^h*                         S r SSKrSSKrSSKJr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Jr  SS
KJrJrJr  SSKJr   " S S5      rg)z
Supports using JWT's for authenticating into the proxy. 

Currently only supports admin. 

JWT token must have 'litellm_proxy_admin' in scope. 
    N)Optionalcast)x509)default_backend)serialization)verbose_proxy_logger)	DualCache)HTTPHandler)JWKKeyValue
JWTKeyItemLiteLLM_JWTAuth)PrismaClientc                      \ rS rSr% Sr\\   \S'   \\S'     S#S jr	 S$S\\   S\S\
S	\SS4
S
 jjrS\4S jrS\S\4S jrS\S\\   S\\   4S jrS\4S jrS\4S jrS\S\\   S\\   4S jrS%S\\   S\4S jjrS\S\\   S\\   4S jrS\S\\   S\\   4S jrS\S\\   S\\   4S jrS\S\4S jrS\\   S\4S jrS\S\\   S\\   4S jrS\S\4S jr S\S\4S  jr!S! r"S"r#g)&
JWTHandler   z
- treat the sub id passed in as the user id
- return an error if id making request doesn't exist in proxy user table
- track spend against the user id
- if role="litellm_proxy_user" -> allow making calls + info. Can not edit budgets
prisma_clientuser_api_key_cachereturnNc                 0    [        5       U l        SU l        g )Nr   )r
   http_handlerleewayselfs    W/home/james-whalen/.local/lib/python3.13/site-packages/litellm/proxy/auth/handle_jwt.py__init__JWTHandler.__init__#   s     (M    litellm_jwtauthr   c                 4    Xl         X l        X0l        X@l        g N)r   r   r   r   )r   r   r   r   r   s        r   update_environmentJWTHandler.update_environment)   s     +"4.r   tokenc                 @    UR                  S5      n[        U5      S:H  $ )N.   )splitlen)r   r#   partss      r   is_jwtJWTHandler.is_jwt5   s    C 5zQr   scopesc                 :    U R                   R                  U;   a  gg)NTF)r   admin_jwt_scope)r   r,   s     r   is_adminJWTHandler.is_admin9   s    //69r   default_valuec                      U R                   R                  b  XR                   R                     nU$ S n U$ ! [         a    Un U$ f = fr    )r   end_user_id_jwt_fieldKeyErrorr   r#   r1   user_ids       r   get_end_user_idJWTHandler.get_end_user_id>   s]    	$##99E 4 4 J JK
     	$#G	$   /8 8 AAc                 4    U R                   R                  c  gg)z@
Returns:
- True: if 'team_id_jwt_field' is set
- False: if not
FT)r   team_id_jwt_fieldr   s    r   is_required_team_idJWTHandler.is_required_team_idJ   s     119r   c                     U R                   R                  b*  [        U R                   R                  [        5      (       a  gg)zh
Returns:
- True: if 'user_allowed_email_domain' is set
- False: if 'user_allowed_email_domain' is None
TF)r   user_allowed_email_domain
isinstancestrr   s    r   is_enforced_email_domain#JWTHandler.is_enforced_email_domainT   s=     99E*  ::CK
 K
 r   c                      U R                   R                  b  XR                   R                     nU$ U R                   R                  b  U R                   R                  nU$ S n U$ ! [         a    Un U$ f = fr    )r   r;   team_id_defaultr4   )r   r#   r1   team_ids       r   get_team_idJWTHandler.get_team_ida   s    	$##55A 4 4 F FG  %%55A..>>
     	$#G	$s   /A' -A' "A' 'A76A7valid_user_emailc                 :    USL a  gU R                   R                  $ )z_
Returns:
- True: if 'user_id_upsert' is set AND valid_user_email is not False
- False: if not
F)r   user_id_upsert)r   rI   s     r   is_upsert_user_idJWTHandler.is_upsert_user_idm   s!     u$##222r   c                      U R                   R                  b  XR                   R                     nU$ Un U$ ! [         a    Un U$ f = fr    )r   user_id_jwt_fieldr4   r5   s       r   get_user_idJWTHandler.get_user_idw   s]    	$##55A 4 4 F FG
  (   	$#G	$r9   c                      U R                   R                  b  XR                   R                     nU$ S n U$ ! [         a    Un U$ f = fr    )r   user_email_jwt_fieldr4   )r   r#   r1   
user_emails       r   get_user_emailJWTHandler.get_user_email   s`    	'##88D"#7#7#L#LM

  "
   	'&J	'r9   c                      U R                   R                  b  XR                   R                     nU$ S n U$ ! [         a    Un U$ f = fr    )r   org_id_jwt_fieldr4   )r   r#   r1   org_ids       r   
get_org_idJWTHandler.get_org_id   s]    	###44@33DDE
     	#"F	#r9   c                      [        US   [        5      (       a  US   R                  5       nU$ [        US   [        5      (       a  US   nU$ [	        S[        US   5       S35      e! [         a    / n U$ f = f)NscopezUnmapped scope type - z. Supported types - list, str.)r@   rA   r'   list	Exceptiontyper4   )r   r#   r,   s      r   
get_scopesJWTHandler.get_scopes   s    	%.#..w--/  E'ND11w   ,T%.-A,BB`a   	F	s   +A) A) A) )A98A9kidc                 @  #    [         R                  " S5      nUc  [        S5      eU R                  R	                  S5      I S h  vN nUc  U R
                  R                  U5      I S h  vN nUR                  5       nSU;   a  UR                  5       S   nOUnU R                  R                  SUU R                  R                  S9I S h  vN   OUnU R                  XaS9nUc   [        SU SU S	U S
[        U5       35      e[        [        U5      $  N N NM7f)NJWT_PUBLIC_KEY_URLz,Missing JWT Public Key URL from environment.litellm_jwt_auth_keyskeys)keyvaluettl)rg   rc   z"No matching public key found. kid=z, keys_url=z, cached_keys=z, len(keys)=)osgetenvr_   r   async_get_cacher   getjsonasync_set_cacher   public_key_ttl
parse_keysr(   r   dict)r   rc   keys_urlcached_keysresponseresponse_jsonrg   
public_keys           r   get_public_keyJWTHandler.get_public_key   s;    9912JKK 33CC#
 
 !..228<<H$MMOM&$,MMOF$;$))99+((77 :    D__$_8
4SEXJn]h\iiuvyz~v  vA  B  D*%%3
 =s8   ADD%D+D,A"DDA
DDDrg   c                    S n[        U5      S:X  aj  [        U[        5      (       a  UR                  SS 5      U:X  d  Uc  UnU$ [        U[        5      (       a!  US   R                  SS 5      U:X  d  Uc  US   nU$ [        U5      S:  a]  U HW  n[        U[        5      (       a  UR                  SS 5      nOS nUc  M2  [        U[        5      (       d  MI  Uc  MN  XR:X  d  MU  UnMY     U$ )N   rc   r   )r(   r@   rs   rn   r^   )r   rg   rc   rx   rh   key_kids         r   rr   JWTHandler.parse_keys   s    +/
t9>$%%488E4+@C+G3;!
& % D$''QE4(C/3;!!W
  Y]c4((!ggeT2G"GO"3--+!$J  r   rT   c                     U R                   R                  c  gUR                  S5      S   nX R                   R                  :X  a  gg)NT@F)r   r?   r'   )r   rT   email_domains      r   is_allowed_domainJWTHandler.is_allowed_domain   sD    99A!'',R0//IIIr   c           	        #    / SQn[         R                  " S5      nS nUc  SS0nSS KnSSKJn  UR                  U5      n[        R                  " SU5        UR                  SS 5      nU R                  US	9I S h  vN n	U	b  [        U	[        5      (       a  0 n
S
U	;   a  U	S
   U
S
'   SU	;   a  U	S   U
S'   SU	;   a  U	S   U
S'   SU	;   a  U	S   U
S'   UR                  [        R                  " U
5      5      n UR                  UUUUUU R                   S9nU$ U	b  [        U	[&        5      (       a   [(        R*                  " U	R-                  5       [/        5       5      nUR1                  5       R3                  [4        R6                  R8                  [4        R:                  R<                  5      nUR                  UUUUUS9nU$ [%        S5      e GNS! UR"                   a    [%        S5      e[$         a  n[%        S['        U5       35      eS nAff = f! UR"                   a    [%        S5      e[$         a  n[%        S['        U5       35      eS nAff = f7f)N)RS256RS384RS512PS256PS384PS512JWT_AUDIENCE
verify_audFr   )RSAAlgorithmz
header: %src   )rc   ktyne)
algorithmsoptionsaudiencer   zToken ExpiredzValidation fails: )r   r   r   zInvalid JWT Submitted)rk   rl   jwtjwt.algorithmsr   get_unverified_headerr   debugrn   ry   r@   rs   from_jwkro   dumpsdecoder   ExpiredSignatureErrorr_   rA   r   load_pem_x509_certificateencoder   rx   public_bytesr   EncodingPEMPublicFormatSubjectPublicKeyInfo)r   r#   r   r   decode_optionsr   r   headerrc   rx   jwkpublic_key_rsapayloadr   certrh   s                   r   auth_jwtJWTHandler.auth_jwt   sh     L
99^,*E2N/**51""<8jj%..3.77
!jT&B&BC
"'.E

"'.E
j %c?Cj %c?C)224::c?CN?**")*%;; %   #
:s(C(C?55%%'):
 oo'44!**..!..CC **)%* %   /00y 86 ,, 100 ?"4SVH =>>?0 ,, 100 ?"4SVH =>>?s\   A:I<G=A;I9G I2BH I$H6HHI$I9IIIc                 T   #    U R                   R                  5       I S h  vN   g  N7fr    )r   closer   s    r   r   JWTHandler.close<  s     %%'''s   (&()r   r   r   r   r   )r   N)r   r    )$__name__
__module____qualname____firstlineno____doc__r   r   __annotations__r	   r   r   intr!   rA   r*   r^   boolr/   rs   r7   r<   rB   rG   rL   rP   rU   rZ   ra   ry   r   r   rr   r   r   r   __static_attributes__ r   r   r   r      s    L))!!	 
-
 &
 )	

 
 

C t  


*23-
	#
T $ 
 
hsm 
QT 
3(4. 3D 3 hsm QT 

*23-
	#
 Xc] xPS}   & &$ &B{ # 8JCW 2C D P1C P1D P1d(r   r   )r   ro   rk   typingr   r   cryptographyr   cryptography.hazmat.backendsr   cryptography.hazmat.primitivesr   litellm._loggingr   litellm.caching.cachingr	   'litellm.llms.custom_httpx.httpx_handlerr
   litellm.proxy._typesr   r   r   litellm.proxy.utilsr   r   r   r   r   <module>r      s<     	 !  8 8 1 - ? I I ,e( e(r   