
    g*                         d dl Z d dlZd dlZd dlmZmZ d dlmZmZmZ d dl	m
Z
 d dlZd dlmZ d dlmZmZmZmZmZmZmZmZmZmZ dZ G d d	      Zy)
    N)datetimetimezone)AnyListOptional)status)verbose_proxy_logger)
GenerateKeyRequest
KeyRequestLiteLLM_AuditLogsLiteLLM_VerificationTokenLitellmTableNamesProxyErrorTypesProxyExceptionUpdateKeyRequestUserAPIKeyAuthWebhookEventzlitellm/c                      e Zd Ze	 ddedededee   fd       Z	e	 dde
dedededee   f
d       Ze	 dded	ee   dededee   f
d
       Zededefd       Zededefd       Zed	ee   fd       Zedefd       Zy)KeyManagementEventHooksNdataresponseuser_api_key_dictlitellm_changed_byc                   K   ddl m} ddlm} | j                  du rt
        j                  |       d{    t        j                  du rt        j                  |t              }t        j                   |t        t        t        j                                t#        j$                  t&        j(                        |xs |j*                  xs ||j,                  t.        j0                  |j3                  dd      d	|d
	                   t
        j5                  | j6                  xs dt        j                           |j3                  dd             d{    y7 )7 w)z
        Hook that runs after a successful /key/generate request

        Handles the following:
        - Sending Email with Key Details
        - Storing Audit Logs for key generation
        - Storing Generated Key in DB
        r   create_audit_log_for_updatelitellm_proxy_admin_nameTNdefaulttoken_id created	id
updated_at
changed_bychanged_by_api_key
table_name	object_idactionupdated_valuesbefore_valuerequest_datazvirtual-key-token)secret_namesecret_token)+litellm.proxy.management_helpers.audit_logsr   litellm.proxy.proxy_serverr   send_invite_emailr   _send_key_created_emaillitellmstore_audit_logsjsondumpsstrasynciocreate_taskr   uuiduuid4r   nowr   utcuser_idapi_keyr   KEY_TABLE_NAMEget$_store_virtual_key_in_secret_manager	key_alias)r   r   r   r   r   r   _updated_valuess          c/var/www/openai/venv/lib/python3.12/site-packages/litellm/proxy/hooks/key_management_event_hooks.pyasync_key_generated_hookz0KeyManagementEventHooks.async_key_generated_hook   s$    	
 	H!!T))AA(KKK ##t+"jj3?O+!2tzz|,#+<<#=#5 $4,44$43+<+D+D#4#C#C"*,,z2">('6%)"$ &JJGL*G!gr2 K 
 	
 	
/ L.	
s"   3E#ED"E#E!E#!E#existing_key_rowc                 4  K   ddl m} ddlm} t        j
                  du rt        j                  | j                  d      t              }|j                  d      }t        j                  |t              }t        j                   |t        t        t        j                               t        j                  t         j"                        |xs |j$                  xs ||j&                  t(        j*                  | j,                  d||	      	             y
w)z~
        Post /key/update processing hook

        Handles the following:
        - Storing Audit Logs for key update
        r   r   r   Texclude_noner   updatedr$   r.   N)r3   r   r4   r   r7   r8   r9   r:   r;   r<   r=   r   r>   r?   r   r@   r   rA   rB   rC   r   rD   key)	r   rK   r   r   r   r   r   rH   _before_values	            rI   async_key_updated_hookz.KeyManagementEventHooks.async_key_updated_hookO   s     	
 	H ##t+"jj)EsSO,11t1DM JJ}cBM+!2tzz|,#+<<#=#5 $4,44$43+<+D+D#4#C#C"&((('6%2"" 	s   DDkeys_being_deletedc                   K   ddl m} ddlm}m} t
        j                  du r.| j                  !| j                  D ]  }|j                  |dd       d{   }	|	/t        d	| d
t        j                  dt        j                        |	j                  d      }	t        j                  |	t               }
t#        j$                   |t'        t!        t)        j*                               t-        j.                  t0        j2                        |xs |j4                  xs ||j6                  t8        j:                  |dd|
	                    t<        j?                  |       d{    y7 7 w)z
        Post /key/delete processing hook

        Handles the following:
        - Storing Audit Logs for key deletion
        r   r   )r   prisma_clientTNrP   find_unique)r0   r)   
query_typezKey z
 not found)messagetypeparamcoderM   r   deletedz{}r$   r.   )rS   ) r3   r   r4   r   rU   r7   r8   keysget_datar   r   bad_request_errorr   HTTP_404_NOT_FOUNDr9   r:   r;   r<   r=   r   r>   r?   r   r@   r   rA   rB   rC   r   rD   r   (_delete_virtual_keys_from_secret_manager)r   rS   r   r   r   r   r   rU   rP   key_row_key_rows              rI   async_key_deleted_hookz.KeyManagementEventHooks.async_key_deleted_hook|   sY    	
 	W ##t+		0Eyy - 6 6%M !7 !  ?("&se: 6,>>##66	  ",,D,9::gs;##/%6"4::<0'/||HLL'A'9 (8088(87/@/H/H'8'G'G&)#,+/)1&! !D &NN1 O 
 	
 	
 	IB	
s%   AE6E1DE6+E4,E64E6r1   r2   c                 &  K   t         j                  {t         j                  j                  du r^ddlm} t        t         j                  |      r=t         j                  j                  t        j                  |       |       d{    yyyy7 w)z
        Store a virtual key in the secret manager

        Args:
            secret_name: Name of the virtual key
            secret_token: Value of the virtual key (example: sk-1234)
        NTr   BaseSecretManager)r1   secret_value)
r7   _key_management_settingsstore_virtual_keys+litellm.secret_managers.base_secret_managerrg   
isinstancesecret_manager_clientasync_write_secretr   _get_secret_name)r1   r2   rg   s      rI   rF   z<KeyManagementEventHooks._store_virtual_key_in_secret_manager   s      ++7//BBdJ
 g;;=NO!77JJ$;$L$L'% &2	 K    P K 8s   BBB	Breturnc                     t         j                  j                  j                  d      rt         j                  j                   |  S t         j                  j                   d|  S )N/)r7   ri   prefix_for_stored_virtual_keysendswithr1   s    rI   ro   z(KeyManagementEventHooks._get_secret_name   s`    ++JJSS
 66UUVWbVcdd66UUVVWXcWdee    c                   K   t         j                  t         j                  j                  du rddlm} t        t         j                  |      r}| D ]w  }|j                  Ft         j                  j                  t        j                  |j                               d{    Ut        j                  d|j                   d       y yyyy7 .w)z
        Deletes virtual keys from the secret manager

        Args:
            keys_being_deleted: List of keys being deleted, this is passed down from the /key/delete operation
        NTr   rf   ru   z]KeyManagementEventHooks._delete_virtual_key_from_secret_manager: Key alias not found for key z(. Skipping deletion from secret manager.)r7   ri   rj   rk   rg   rl   rm   rG   async_delete_secretr   ro   r	   warningr0   )rS   rg   rP   s      rI   ra   z@KeyManagementEventHooks._delete_virtual_keys_from_secret_manager   s      ++7//BBdJ g;;=NO1==4")"?"?"S"S,C,T,T$'MM-" #T #   188"  AD  AJ  AJ  @K  Ks  !t  2 P K 8s   BC C!/Cc                   K   ddl m}m} d|j                  dg       vrt	        d      t        ddd| j                  d	d
      | j                  dd      | j                  dd      | j                  dd       | j                  dd      | j                  dd       	      }t        j                  |j                  j                  |             y w)Nr   )general_settingsproxy_logging_objemailalertingzyEmail alerting not setup on config.yaml. Please set `alerting=['email']. 
Docs: https://docs.litellm.ai/docs/proxy/email`key_createdrP   zAPI Key Createdr0   r"   spendg        
max_budgetrB   team_idzDefault TeamrG   )	eventevent_groupevent_messager0   r   r   rB   r   rG   )webhook_event)
r4   r{   r|   rE   
ValueErrorr   r<   r=   slack_alerting_instance&send_key_created_or_user_invited_email)r   r{   r|   r   s       rI   r6   z/KeyManagementEventHooks._send_key_created_email   s     R*..z2>> M  +,,w+,,w,||L#6LLD1LLN;ll;5

 	55\\# ] 	
s   CC)N)__name__
__module____qualname__staticmethodr
   dictr   r   r;   rJ   r   r   rR   r   r   r   rd   rF   ro   ra   r6    rv   rI   r   r      sf   
 -1	.
 .
.
 *.
 %SM	.
 .
`  -1*** * *	*
 %SM* *X  -1;; !:;; ; *	;
 %SM; ;z  SV  . fc fc f f  !:; 8 
 
 
rv   r   )r<   r9   r>   r   r   typingr   r   r   fastapir   r7   litellm._loggingr	   litellm.proxy._typesr
   r   r   r   r   r   r   r   r   r   "LITELLM_PREFIX_STORED_VIRTUAL_KEYSr   r   rv   rI   <module>r      sC       ' & &   1   &0 "u
 u
rv   