hl`dZddlZddlZddlZddlmZmZmZmZm Z m Z m Z ddl m Z ddlmZddlmZddlmZddlmZdd lmZdd lmZdd lmZdd lmZmZmZm Z m!Z!m"Z"m#Z#m$Z$m%Z%m&Z&m'Z'm(Z(m)Z)m*Z*m+Z+m,Z,m-Z-dd l.m/Z/ddl0m1Z1m2Z2ddl3m4Z4m5Z5m6Z6m7Z7m8Z8m9Z9m:Z:m;Z;my)z Supports using JWT's for authenticating into the proxy. Currently only supports admin. JWT token must have 'litellm_proxy_admin' in scope. N)AnyListLiteralOptionalSetTuplecast)x509)default_backend) serialization HTTPException)verbose_proxy_logger) DualCache)get_nested_value) HTTPHandler) RBAC_ROLES JWKKeyValueJWTAuthBuilderResult JWTKeyItemLiteLLM_EndUserTableLiteLLM_JWTAuthLiteLLM_OrganizationTableLiteLLM_TeamTableLiteLLM_UserTableLitellmUserRolesMemberProxyErrorTypesProxyException ScopeMappingSpanTeamMemberAddRequestUserAPIKeyAuth)can_team_access_model) PrismaClient ProxyLogging) _allowed_routes_checkallowed_routes_checkget_actual_routesget_end_user_objectget_org_objectget_role_based_modelsget_role_based_routesget_team_objectget_user_objectc eZdZUdZeeed<eed< d*dZ d+deedede de ddf d Z e d e fd Zd edeefd Zd edeefd ZdedefdZd edee fdZd edee dee fdZdefdZdefdZd edee dee fdZd,deedefdZd edee dee fdZd edeee deee fdZd edeefdZ d edeee deee fdZ!deee defdZ"d edee dee fdZ#d edee dee fdZ$d edee dee fd Z%d edee fd!Z&d"ee defd#Z'd$e(d"ee dee)fd%Z*d&e defd'Z+d e defd(Z,d)Z-y)- JWTHandlerz - 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_clientuser_api_key_cachereturnNc0t|_d|_y)Nr)r http_handlerleewayselfs Y/var/www/Befach/backend/env/lib/python3.12/site-packages/litellm/proxy/auth/handle_jwt.py__init__zJWTHandler.__init__Es(M litellm_jwtauthr8c<||_||_||_||_yN)r3r4r>r8)r:r3r4r>r8s r;update_environmentzJWTHandler.update_environmentKs#+"4. r=tokenc@|jd}t|dk(S)N.)splitlen)rBpartss r;is_jwtzJWTHandler.is_jwtWs C 5zQr=c|jjy|j|d}|syt|}|jjD]}|j|vs|j cSy)a Returns the RBAC role the token 'belongs' to based on role mappings. Args: token (dict): The JWT token containing role information Returns: Optional[RBAC_ROLES]: The mapped internal RBAC role if a mapping exists, None otherwise Note: The function handles both single string roles and lists of roles from the JWT. If multiple mappings match the JWT roles, the first matching mapping is returned. NrB default_value)r> role_mappings get_jwt_rolesetrole internal_role)r:rBjwt_role jwt_role_set role_mappings r;_rbac_role_from_role_mappingz'JWTHandler._rbac_role_from_role_mapping\sr    - - 5$$5$E8}  00>>L  L0#111? r=c|j|}|j|}|j|d}|rtjS|j |dtj S|j|dtjS|"|j|rtjS|j|x}r|Sy)a^ Returns the RBAC role the token 'belongs' to. RBAC roles allowed to make requests: - PROXY_ADMIN: can make requests to all routes - TEAM: can make requests to routes associated with a team - INTERNAL_USER: can make requests to routes associated with a user Resolves: https://github.com/BerriAI/litellm/issues/6793 Returns: - PROXY_ADMIN: if token is admin - TEAM: if token is associated with a team - INTERNAL_USER: if token is associated with a user - None: if token is not associated with a team or user rBscopesNrK) user_roles) get_scopesis_adminget_user_rolesr PROXY_ADMIN get_team_idTEAM get_user_id INTERNAL_USERis_allowed_user_rolerU)r:rBrYr\rZ rbac_roles r; get_rbac_rolezJWTHandler.get_rbac_role{s"u-===/((uD(I #// /   E  > J#(( (   E  > J#11 1  #(A(A!)B) $11 1;;%;H HY H r=rYc6|jj|vryy)NTF)r>admin_jwt_scope)r:rYs r;r\zJWTHandler.is_admins    / /6 9r=c|jj>|j|jj||jjSgSr@)r>team_ids_jwt_fieldget)r:rBs r;get_team_ids_from_jwtz JWTHandler.get_team_ids_from_jwtsM  3 3 ? $..AABN--@@A A r=rLc |jj||jj}|Sd} |S#t$r|}Y|SwxYwr@)r>end_user_id_jwt_fieldKeyErrorr:rBrLuser_ids r;get_end_user_idzJWTHandler.get_end_user_ids_ $##99E 4 4 J JK   $#G $/88 AAc2|jjyy)z` Returns: - True: if 'team_id_jwt_field' is set - False: if not FT)r>team_id_jwt_fieldr9s r;is_required_team_idzJWTHandler.is_required_team_ids    1 1 9r=cz|jj%t|jjtryy)z Returns: - True: if 'user_allowed_email_domain' is set - False: if 'user_allowed_email_domain' is None TF)r>user_allowed_email_domain isinstancestrr9s r;is_enforced_email_domainz#JWTHandler.is_enforced_email_domains7    9 9 E*  : :CK r=c |jj||jj}|S|jj|jj}|Sd} |S#t$r|}Y|SwxYwr@)r>rtteam_id_defaultrn)r:rBrLteam_ids r;r_zJWTHandler.get_team_ids $##55A 4 4 F FG %%55A..>>  $#G $s/A&,A&!A&& A54A5valid_user_emailc8|dury|jjS)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:r~s r;is_upsert_user_idzJWTHandler.is_upsert_user_ids! u $##222r=c |jj||jj}|S|} |S#t$r|}Y|SwxYwr@)r>user_id_jwt_fieldrnros r;razJWTHandler.get_user_ids_ $##55A 4 4 F FG ( $#G $rrc |jj$t||jj|}|S|} |S#t$r|}Y|SwxYw)zn Returns the user role from the token. Set via 'user_roles_jwt_field' in the config. datakey_pathdefault)r>user_roles_jwt_fieldrrnr:rBrLrZs r;r]zJWTHandler.get_user_rolessk '##88D-!11FF) +  '&J '8AA AAc|jjsy|j|g}|sy|jjD]9}|D]2}tj||js$|j ccS;y)z(Map roles from JWT to LiteLLM user rolesNrK)r>jwt_litellm_role_maprNfnmatchrR litellm_role)r:rB jwt_rolesmappingrPs r;map_jwt_role_to_litellm_rolez'JWTHandler.map_jwt_role_to_litellm_rolesu##88%%E%D ++@@G!??4)9)9:"///"Ar=c |jj$t||jj|}|S|} |S#t$r|}Y|SwxYw)z Generic implementation of `get_user_roles` that can be used for both user and team roles. Returns the jwt role from the token. Set via 'roles_jwt_field' in the config. r)r>roles_jwt_fieldrrnrs r;rNzJWTHandler.get_jwt_rolesk '##33?-!11AA) +  '&J 'rrZc`|+jjtfd|Dryy)zl Returns the user role from the token. Set via 'user_allowed_roles' in the config. c3NK|]}|jjvywr@)r>user_allowed_roles).0rPr:s r; z2JWTHandler.is_allowed_user_role..6s'DH,,???s"%TF)r>rany)r:rZs` r;rczJWTHandler.is_allowed_user_role-s:  "$$77CLVr=c |jj||jj}|Sd} |S#t$r|}Y|SwxYwr@)r>user_email_jwt_fieldrn)r:rBrL user_emails r;get_user_emailzJWTHandler.get_user_email=sb '##88D"4#7#7#L#LM "  '&J 'rrc |jj||jj}|S|} |S#t$r|}Y|SwxYwr@)r>object_id_jwt_fieldrn)r:rBrL object_ids r; get_object_idzJWTHandler.get_object_idIsb &##77C!$"6"6"J"JK *  &%I &rrc |jj||jj}|Sd} |S#t$r|}Y|SwxYwr@)r>org_id_jwt_fieldrn)r:rBrLorg_ids r; get_org_idzJWTHandler.get_org_idSs_ ###44@t33DDE    #"F  #rrc t|dtr|dj}|St|dtr|d}|St dt |dd#t $rg}Y|SwxYw)NscopezUnmapped scope type - z. Supported types - list, str.)rxryrFlist Exceptiontypern)r:rBrYs r;r[zJWTHandler.get_scopes]s %.#.w--/ E'ND1w  ,T%.-A,BB`a F  s&AAA A.-A.kidcKtjd}| td|jdDcgc]}|j }}|D]}d|}|j j |d{}||jj|d{}|j} d| vr|jd} n| } |j j|| |jjd{n|} |j| |} | tt| cStd|d |cc}w777Jw) NJWT_PUBLIC_KEY_URLz,Missing JWT Public Key URL from environment.,litellm_jwt_auth_keys_keys)keyvaluettl)rrz#No matching public key found. keys=z, kid=)osgetenvrrFstripr4async_get_cacher7rjjsonasync_set_cacher>public_key_ttl parse_keysr dict) r:rkeys_urlurl keys_url_listkey_url cache_key cached_keysresponse response_jsonr public_keys r;get_public_keyzJWTHandler.get_public_keyls[9912  JK K08s0CDD D$G0 :I $ 7 7 G G RRK"!%!2!2!6!6w!?? (  ]*(0 (?D(D--==!,,;;> #d$%488E4+@C+G3;! &%D$'Q E4(C/3;!!W Y]c4(!ggeT2G"GO"3-+3!$Jr=rc|jjy|jdd}||jjk(ryy)NT@F)r>rwrF)r:r email_domains r;is_allowed_domainzJWTHandler.is_allowed_domainsF    9 9 A!'',R0 4//II Ir=cKgd}tjd}d}|ddi}ddl}ddlm}|j |}t jd||jdd}|j| d{} | t| trzi} d | vr| d | d <d| vr| d| d<d | vr| d | d <d | vr| d | d <|jtj| } |j|| ||||j } | S| t| t&r t)j*| j-t/}|j1j3t4j6j8t4j:j<}|j|||||} | St%d7C#|j"$r t%dt$$r} t%dt'| d} ~ wwxYw#|j"$r t%dt$$r} t%dt'| d} ~ wwxYww)N)RS256RS384RS512PS256PS384PS512 JWT_AUDIENCE verify_audFr) RSAAlgorithmz header: %sr)rktyne) algorithmsoptionsaudiencer8z Token ExpiredzValidation fails: )rrrzInvalid JWT Submitted)rrjwtjwt.algorithmsrget_unverified_headerrdebugrjrrxrfrom_jwkrdumpsdecoder8ExpiredSignatureErrorrryr load_pem_x509_certificateencoder r public_bytesr EncodingPEM PublicFormatSubjectPublicKeyInfo)r:rBrrdecode_optionsrrheaderrrjwkpublic_key_rsapayloadrcertrs r;auth_jwtzJWTHandler.auth_jwtscL 99^,  *E2N/**51""<8jj%..3.77  !jT&BC "'.E  "'.E j %c?Cj %c?C)224::c?CN ?**")*%;; % # :s(C ?55%%'): oo'44!**..!..CC **)%* %/00y86,, 100 ?"4SVH =>> ?0,, 100 ?"4SVH =>> ?s\A:I<F?=A,I*"G I BH3 I"H$G;;HI"I%H<<IIcTK|jjd{y7wr@)r7closer9s r;rzJWTHandler.closes%%'''s (&()r5N)rr@).__name__ __module__ __qualname____doc__rr%__annotations__rr<rintrA staticmethodryrIrrrUrerboolr\rrkrqrurzr_rrar]rrrNrcrrrr[rrrrrrrr=r;r2r2:sL))!!   - & )     c$8J;O>"4"HZ,@"Ht 4DI  *23-  # T $   hsm QT 3(4.3D3hsmQT *249*= $s) * $ 8DT;U *249*= $s) .xS /Bt  *23-  # 4 (SV-Xc]xPS}  c $  $ $$ L{#8JCW2CDP1CP1DP1d(r=r2c|eZdZdZedededededfdZ ededede ededfd Z ed e e d e ed ededd f dZedededed edede edd fdZeded edede ede edede efdZededede edede ededee ee effdZedededeefdZedeede ededede edede ededee ee effdZedededee ee ee effdZede ed e ede ed!e ed"e edede edede ededee e e e!e e"ffd#Z#ede ed$e ed%ed&ededf d'Z$ed(e e d)e efd*Z%ededed(e e de edd f d+Z&ededed ededede edede ededefd,Z'y )-JWTAuthManagerz7Manages JWT authentication and authorization operationsrdgeneral_settingsrouter5Tc rt||}||yt||}|stdd|d|d|y) zU Checks if user is allowed to access the route, based on their role. rdrT) user_routeallowed_routesRole=z not allowed to call route=z. Allowed routes= status_codedetail)r.r(r)rdrrrole_based_routes is_alloweds r;can_rbac_role_call_routez'JWTAuthManager.can_rbac_role_call_routesh22B   $ *, yk)DUGK\]n\op  r=modelc \t||}||y||vrtdd|d|d|y)zU Checks if user is allowed to access the model, based on their role. rTrr z not allowed to call model=z. Allowed models=r )r-r)rdrrrole_based_modelss r;can_rbac_role_call_modelz'JWTAuthManager.can_rbac_role_call_model+s\22B   $  ) )yk)DUGK\]n\op  r=scope_mappingsrY request_dataNc|syg}|D]9}|j|vs|js|j|j;|jd}|sy||vrt dddj ||iy)zE Check if scope allows access to the requested model Nrrerrorz'model={} not allowed. Allowed_models={}r )rmodelsextendrjrformat)rrYrrallowed_modelssmrequested_models r;check_scope_based_accessz'JWTAuthManager.check_scope_based_accessBs Bxx6!bii%%bii0!'**73 . 0FMM' r= jwt_handlerjwt_valid_tokencK|jjdurO| tddtj |||j dtj |||yyw) z/Validate RBAC role and model access permissionsTNrzjUnmatched token passed in. enforce_rbac is set to True. Token must belong to a proxy admin, team, or user.r r)rdrr)rdrr)r> enforce_rbacrrrrjr)rr rrrrds r;check_rbac_rolezJWTAuthManager.check_rbac_roleds  & & 3 3t ; # #H  3 3#!1"&&w/ 4   3 3#!1 4  admin_allowed_routesr*rr) rrYrrprr$rr actual_routess r;check_admin_accessz!JWTAuthManager.check_admin_accesss##6#2)&22 + ; ; (3(C(C(X(XN-^LM@GXYfXgh $   sBBr3r4parent_otel_spanproxy_logging_objc K|j|d}|s5|jdur#td|jjdd}|r-t ||||||jj d{}||fS7w)z"Find and validate specific team IDNrKTz2No team id found in token. Checked team_id field '')r}r3r4r2r3team_id_upsert)r_rurr>rtr/r6)rr r3r4r2r3individual_team_idr*s r;"find_and_validate_specific_team_idz1JWTAuthManager.find_and_validate_specific_team_ids)44!5 "k&E&E&G4&OD[E`E`ErErDsstu  48  /*+#5!1"3*::II !K";..sA7B9B: Bc@|j|}t|}|S)z8Get combined team IDs from groups and individual team_idrW)rkrO)rr team_ids_from_groups all_team_idss r;get_all_team_idszJWTAuthManager.get_all_team_idss) +@@@W/0 r=team_idsrcK|s$|jjr tddy|D]} t|||||d{} | rg| j[| j} t | t r?|rt|| ddr.ttj||j} | r|| fcS|rtdd |d |d y7#t$rYwxYww) z2Find first team with access to the requested modelrzgNo teams found in token. `enforce_team_based_model_access` is set to True. Token must belong to a team.r )NN)r}r3r4r2r3N)rr* llm_routerteam_model_aliasesr&z+No team has access to the requested model: z. Checked teams=z.. Check `/models` to see all available models.) r>enforce_team_based_model_accessrr/rrxrr$r)rr`r) r=rrrr3r4r2r3r}r* team_modelsrs r;find_team_with_model_accessz*JWTAuthManager.find_team_with_model_accesss!**JJ# #EG $3#"/'9%5&7 % ;#5#5#A"-"4"4K!+t4+0"1(3'+/3 &:&6&;&;',0;0K0K& &#*K#777 > D_DUUefneop^_  G4  s;,CC C A*C -C C  CCCCcK|j|d}d}|jr|dn|j|}|j||}|||fSw)z$Get user email and validation statusNrKF)r)rrzrra)rr rr~rps r; get_user_infozJWTAuthManager.get_user_info s !//!0    / / 1% 22j2I  ))!*  $444sAArr.r~c Kd} |r|rt||||| d{nd} d} |r/|r+t||||j||| ||d{nd} d} |r|rt||||| d{nd} | | | fS7^7-7w)z#Get user, org, and end user objectsN)rr3r4r2r3)r~)rpr3r4rr2r3r sso_user_id)r.r3r4r2r3)r,r0rr+) rprrr.r~rr3r4r2r3r-r+r,s r; get_objectszJWTAuthManager.get_objects s$;? %!"/'9%5&7  48 &#"/'9#.#@#@)9$A$&6&7) '    ";? * +"/'9%5&7  J77Y &s3BA;2BA=B/A?0 B=B?Br}r"r)c.|r|s|s|s tddy)zUIf enforce_rbac is true, validate that a valid rbac id is returned for spend trackingrzrNo user or team id found in token. enforce_rbac is set to True. Token must belong to a proxy admin, team, or user.r Tr rpr}r"r)s r;validate_object_idz!JWTAuthManager.validate_object_idcs( wwL r=r+r*cnKddlm}|sy|sy|jD]*}|js|j|jk(s*yt t |jd|j } ||ttjd{tjd |jd |j y74#t$rX}|jtjk(r4tjd |jd |j Yd}~y|d}~wwxYww) z Map user to teams. - If user is not in team, add them to the team - If user is in team, do nothing r)team_member_addNuser)rprP)memberr})r')ruser_api_key_dictzSuccessfully added user z to team zUser z is already a member of team )1litellm.proxy.management_endpoints.team_endpointsrMmembers_with_rolesrpr"rr}r#rr^rrrrrteam_member_already_in_team)r+r*rMrOrrs r;map_user_to_teamsz JWTAuthManager.map_user_to_teamsrs- V"44F~~&..K4G4G"G5$#++ ''   !"0.::#   ! & &*;+>+>*?yI\I\H]^ #  vvDDD$**K//00MkNaNaMbc sS(D5D5/D55%CC3CD5C D2A D-&D5+D--D22D5c`K|jjsy||y|j|}|rq|j|jk7rX|j j jd|jid|jid{|j|_t|j|}t|jxsg}||z }||z }|s|rEddl m } | |jt|t|d{t||_ y77w)ak Sync user role and team memberships with JWT claims The goal of this method is to ensure: 1. The user role on LiteLLM DB is in sync with the IDP provider role 2. The user is a member of the teams specified in the JWT token This method is only called if sync_user_role_and_teams is set to True in the JWT config. Nrpr')whererr)patch_team_membership)rpteams_ids_to_add_user_toteams_ids_to_remove_user_from)r>sync_user_role_and_teamsrr'rdblitellm_usertableupdaterprOrkteams/litellm.proxy.management_endpoints.scim.scim_v2rWr) rr r+r3new_role jwt_team_idsexisting_teams teams_to_addteams_to_removerWs r;rZz'JWTAuthManager.sync_user_role_and_teamss3 **CC  -"7;;OL  --?""44;; +"5"56!8>>2<  %-NNK !;<<_MN [..4"5#n4 (<7 ? (#++)-l);.2?.C   !%\ 2K -   s%B D. D*BD.D,D.,D.c bK|j|d{} |jjr(|jj| s tdd|j | } t j || |||| d{|j| } |jjrC|jjr-t j|jj| |||j| d} |j| } t j|| d{\} }}|j| d}|j| d}d}d}|j| d} | r-| r+| tj k(r| }n| tj"k(r| } t j%|| || ||d{}|r|S|s%t j'|| ||||d{\}}|sO|sMt j)|| }t j+||j-d|||||| d{\}}t j/| ||||||||| d{\}}}t j1|| || d{t j3|| d{t j5| ||j-d dd|r |j6tj8k(rd}nd}t;|||| |||||| S7777W727777{w)z-Main authentication and authorization builderrWNrzInvalid JWT tokenr )rrYrrrKr)r=rrrr3r4r2r3) rprrr.r~rr3r4r2r3)rr r+r3)r+r*r"FrJT) r)r}r*rpr+rr-r.r,rB)rr>custom_validaterrerr#r[enforce_scope_based_accessrrrrErrqrr`rbr1r8r<rCrjrHrZrTrKr'r^r)r$rrrrr3r4r2r3r rdrYrrprr~rr.r}r* admin_resultr;r+r-r,r)s r; auth_builderzJWTAuthManager.auth_builders'2&:&:&:&I I  & & 6 6..>>O# #.  --O-D ,,          ''o'>  ' ' B B++::  3 3*::II)!1 4  --OSW-X ''o'>6D6R6R 7 1 -- ''oT'R!11!2 "&37 --OSW-X ,111#.<<<#,>>      %GG" !  7)::;XL)7)S)S% , 0 0 9'+#5!1"3*T * $ G[:H9S9S!#-#'1-/:T : 4 0 Z55#+#' 6   ..##/   )))--neD *  ;004D4P4PP!N"N#)##!#+  O!J 41 (  $ 4   sL/LA1L/ L B-L/8L9BL/ L! 'L/2L$3AL/L')L/-L).$L/L+L/1L-2A'L/L/L/!L/$L/'L/)L/+L/-L/)(rrrrrrrryrrrrrr rr2r#rrr1r%rr!r&rrr8rr<rCrrErrrrHrKrTrZrirr=r;rr s'A  8}  ,\*S    B    J'  4# # # # # #  #  #  & '# # J/// -/& / #4. / ( / x}h'899 ://@j4CPSH9c(9!#99 9  - 9 & 9#4.9(9 x}h'899 :99v555 x}hsmXd^; <55(@8#@8SM@8 @8c] @8 #4. @8  @8 -@8&@8#4.@8(@8 "#*+%& ( @8@8D # #      0/00/000d.../0. - .  ..`^ ^ ^ ^  ^  ^  - ^ &^ #4.^ (^  ^ ^ r=r)?rrrrtypingrrrrrrr cryptographyr cryptography.hazmat.backendsr cryptography.hazmat.primitivesr fastapirlitellm._loggingrlitellm.caching.cachingr0litellm.litellm_core_utils.dot_notation_indexingr'litellm.llms.custom_httpx.httpx_handlerrlitellm.proxy._typesrrrrrrrrrrrrrr r!r"r#litellm.proxy.auth.auth_checksr$litellm.proxy.utilsr%r& auth_checksr(r)r*r+r,r-r.r/r0r2rrr=r;rws AAA88!1-M?&A:   N(N(bj  j  r=