a
    a*j0                     @   s  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ZddlmZ ddl	m
Z
m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 ddlmZ ddlmZ dd	lmZmZ dd
lmZm Z  ddl!m"Z" ddl#m$Z$ ddl%m&Z& ddl'm(Z(m)Z) ddl*m+Z+ ddl,m-Z- ddl.m/Z/m0Z0 edZ1ee2ede2f f Z3G dd dZ4G dd dZ5dZ6e&e-e)dddZ7G dd dZ8dS )z0
The starlette extension to rate-limit requests
    N)datetime)
formatdateparsedate_to_datetime)wraps)	AnyCallableDictListOptionalSetTupleTypeVarUnion)RateLimitItem)ConfigurationError)MemoryStoragestorage_from_string)
STRATEGIESRateLimiter)Config)MutableHeaders)Request)JSONResponseResponse)Literal   )RateLimitExceeded)Limit
LimitGroupT.c                   @   sP   e Zd ZdZdZdZdZdZdZdZ	dZ
d	Zd
ZdZdZdZdZdZdZdZdS )CZRATELIMIT_ENABLEDZRATELIMIT_HEADERS_ENABLEDZRATELIMIT_STORAGE_URLZRATELIMIT_STORAGE_OPTIONSZRATELIMIT_STRATEGYZRATELIMIT_GLOBALZRATELIMIT_DEFAULTZRATELIMIT_APPLICATIONZRATELIMIT_HEADER_LIMITZRATELIMIT_HEADER_REMAININGZRATELIMIT_HEADER_RESETZRATELIMIT_SWALLOW_ERRORSZRATELIMIT_IN_MEMORY_FALLBACKZ$RATELIMIT_IN_MEMORY_FALLBACK_ENABLEDZRATELIMIT_HEADER_RETRY_AFTERZ"RATELIMIT_HEADER_RETRY_AFTER_VALUEZRATELIMIT_KEY_PREFIXN)__name__
__module____qualname__ENABLEDHEADERS_ENABLEDSTORAGE_URLSTORAGE_OPTIONSSTRATEGYZGLOBAL_LIMITSDEFAULT_LIMITSAPPLICATION_LIMITSHEADER_LIMITHEADER_REMAININGHEADER_RESETSWALLOW_ERRORSIN_MEMORY_FALLBACKIN_MEMORY_FALLBACK_ENABLEDHEADER_RETRY_AFTERHEADER_RETRY_AFTER_VALUE
KEY_PREFIX r4   r4   M/var/www/html/assistant/venv/lib/python3.9/site-packages/slowapi/extension.pyr    -   s"   r    c                   @   s   e Zd ZdZdZdZdZdS )HEADERSr            N)r!   r"   r#   RESET	REMAININGLIMITRETRY_AFTERr4   r4   r4   r5   r6   A   s   r6      )requestexcreturnc                 C   s2   t dd|j idd}| jjj|| jj}|S )z
    Build a simple JSON response that includes the details of the rate limit
    that was hit. If no limit is hit, the countdown is added to headers.
    errorzRate limit exceeded: i  )status_code)r   detailappstatelimiter_inject_headersview_rate_limit)r?   r@   responser4   r4   r5   _rate_limit_exceeded_handlerK   s    
rK   c                   @   sl  e Zd ZdZg g dddi ddg ddddddfedef ee ee ee	e e	e e
eef eeee ee	e eee	e ed dd	d
dZddddZd2eeedddZedddZddddZeedddZeeeee f edddZeeeee f edddZeeee dddd Zedd!d"Zd3ee	ede f  edd#d$d%Z!d4ee	edef  ee	e ee	ee  e	e e	edef  e"eedef f eede f d'd(d)Z#d5ee	edef  ee	ee  e	e e	edef  e"eedef f eed*	d+d,Z$d6eee	edef  e	e e	edef  e"eedef f eed-d.d/Z%d0d1 Z&dS )7Limitera1  
    Initializes the slowapi rate limiter.

    ** parameter **

    * **app**: `Starlette/FastAPI` instance to initialize the extension
     with.

    * **default_limits**: a variable list of strings or callables returning strings denoting global
     limits to apply to all routes. `ratelimit-string` for  more details.

    * **application_limits**: a variable list of strings or callables returning strings for limits that
     are applied to the entire application (i.e a shared limit for all routes)

    * **key_func**: a callable that returns the domain to rate limit by.

    * **headers_enabled**: whether ``X-RateLimit`` response headers are written.

    * **strategy:** the strategy to use. refer to `ratelimit-strategy`

    * **storage_uri**: the storage location. refer to `ratelimit-conf`

    * **storage_options**: kwargs to pass to the storage implementation upon
      instantiation.
    * **auto_check**: whether to automatically check the rate limit in the before_request
     chain of the application. default ``True``
    * **swallow_errors**: whether to swallow errors when hitting a rate limit.
     An exception will still be logged. default ``False``
    * **in_memory_fallback**: a variable list of strings or callables returning strings denoting fallback
     limits to apply when the storage is down.
    * **in_memory_fallback_enabled**: simply falls back to in memory storage
     when the main storage is down and inherits the original limits.
    * **key_prefix**: prefix prepended to rate limiter keys.
    * **enabled**: set to False to deactivate the limiter (default: True)
    * **config_filename**: name of the config file for Starlette from which to load settings
     for the rate limiter. Defaults to ".env".
    * **key_style**: set to "url" to use the url, "endpoint" to use the view_func
    FNT url.)endpointrN   )key_funcdefault_limitsapplication_limitsheaders_enabledstrategystorage_uristorage_options
auto_checkswallow_errorsin_memory_fallbackin_memory_fallback_enabledretry_after
key_prefixenabledconfig_filename	key_stylerA   c                 C   s  t d| _tjd}t|r*|du r*dn|| _|| _g | _	g | _
g | _|pXt|
dk| _t | _g | _|| _i | _|| _|| _|| _|| _|| _|	| _|| _|| _|| _t|D ](}| j	t|| jddddddd	g q|D ](}| j
t|| jddddddd	g q|
D ]*}| jt|| jddddddd	g qi | _i | _d| _ d| _!d| _"t## | _$i | _%G dd	 d	t j&}| j'|  | (t)j*| j| _| (t)j+| j| _| jp| (t)j,d| _| j-| (t)j.i  t/| jp| (t)j0d
fi | j| _1| jp| (t)j2d}|t3vr.t4d| t3| | j1| _5| j-t6j7| j8t6j7| (t)j9dt6j:| j8t6j:| (t)j;dt6j<| j8t6j<| (t)j=dt6j>| j8t6j>| (t)j?di | jp| (t)j@| _| jp| (t)jA| _| (t)jBd}| j
s |r t|| jddddddd	g| _
| (t)jCd}| j	sZ|rZt|| jddddddd	g| _	| (t)jDd}| (t)jEd}| js|rt|| jddddddd	g| _| js|pt| jdk| _| jrtF | _Gt3| | jG| _!dS )z9
        Configure the rate limiter at app level
        Zslowapiz.envNr   Fr   globalc                   @   s   e Zd Zdd ZdS )z*Limiter.__init__.<locals>.BlackHoleHandlerc                  W   s   d S Nr4   )_r4   r4   r5   emit   s    z/Limiter.__init__.<locals>.BlackHoleHandler.emitN)r!   r"   r#   rc   r4   r4   r4   r5   BlackHoleHandler   s   rd   z	memory://zfixed-windowz!Invalid rate limiting strategy %szX-RateLimit-ResetzX-RateLimit-RemainingzX-RateLimit-LimitRetry-After)Hlogging	getLoggerloggerospathisfiler   
app_configr]   _default_limits_application_limits_in_memory_fallbacklen_in_memory_fallback_enabledset_exempt_routes_request_filters_headers_enabled_header_mapping_retry_afterZ	_strategyZ_storage_uriZ_storage_options_auto_check_swallow_errors	_key_func_key_prefix
_key_styleextendr   _route_limits_dynamic_route_limits_storage_dead_fallback_limiter_Limiter__check_backend_counttime_Limiter__last_check_backend_Limiter__marked_for_limitingStreamHandler
addHandlerget_app_configr    r$   r.   r%   updater'   r   r&   _storager(   r   r   _limiterr6   r:   getr-   r;   r,   r<   r+   r=   r1   r2   r3   r*   r)   r0   r/   r   Z_fallback_storage)selfrP   rQ   rR   rS   rT   rU   rV   rW   rX   rY   rZ   r[   r\   r]   r^   r_   Zdotenv_file_existslimitrd   Z
app_limitsZconf_limitsZfallback_enabledZfallback_limitsr4   r4   r5   __init__   s2   





zLimiter.__init__)rA   c                 C   s   | t j_t tt dS )z_
        Starlette startup event handler that links the app with the Limiter instance.
        N)rE   rF   rG   Zadd_exception_handlerr   rK   r   r4   r4   r5   slowapi_startupH  s    zLimiter.slowapi_startup)keydefault_valuerA   c                 C   s&   |r| j ||t|dS | j ||dS )zQ
        Place holder until we find a better way to load config from app
        )defaultcast)r   )rl   type)r   r   r   r4   r4   r5   r   O  s    zLimiter.get_app_configc                 C   sJ   | j tkrd| _ t | j td| j krFt | _|  j d7  _ dS dS )Nr   r7   r   TF)r   MAX_BACKEND_CHECKSr   r   powr   r4   r4   r5   Z__should_check_backendY  s    

zLimiter.__should_check_backendc                 C   s>   z| j   | jd W n ty8   | jd Y n0 dS )z?
        resets the storage if it supports being reset
        z-Storage has been reset and all limits clearedz.This storage type does not support being resetN)r   resetrh   infoNotImplementedErrorwarningr   r4   r4   r5   r   b  s
    
zLimiter.resetc                 C   s*   | j r | jr | jsJ d| jS | jS dS )zT
        The backend that keeps track of consumption of endpoints vs limits
        z=Fallback limiter is needed when in memory fallback is enabledN)r   rq   r   r   r   r4   r4   r5   rG   l  s    zLimiter.limiter)rJ   current_limitrA   c                 C   s`  | j r\| jr\|d ur\t|ts,tdz| jj|d g|d R  }d|d  }|j| j	t
j t|d j |j| j	t
j t|d  |j| j	t
j t| |jd}|d urt| ||}| jdkrt|ntt|t  |j| j	t
j < W nR   | jr>| js>| jd d| _| ||}| jrT| jd n Y n0 |S )	NzHparameter `response` must be an instance of starlette.responses.Responser   r   re   	http-dateBRate limit storage unreachable - falling back to in-memory storageT5Failed to update rate limit headers. Swallowing error)r]   ru   
isinstancer   	ExceptionrG   get_window_statsheadersappendrv   r6   r<   stramountr;   r:   r   max_determine_retry_timerw   r   intr   r=   ro   r   rh   r   rH   ry   	exception)r   rJ   r   window_statsreset_inexisting_retry_after_headerr4   r4   r5   rH   y  sX    

zLimiter._inject_headers)r   r   rA   c                 C   s@  | j r<| jr<|dur<z| jj|d g|d R  }d|d  }t|d j|| jtj < t|d || jtj	 < t||| jtj
 < |d}|durt| ||}| jdkrt|ntt|t  || jtj < W nZ ty:   | jr| js| jd d| _| ||}| jr4| jd n Y n0 |S )	a!  
        Injects 'X-RateLimit-Reset', 'X-RateLimit-Remaining', 'X-RateLimit-Limit'
        and 'Retry-After' headers into :headers parameter if needed.

        Basically the same as _inject_headers, but without access to the Response object.
        -> supports ASGI Middlewares.
        Nr   r   re   r   r   Tr   )r]   ru   rG   r   r   r   rv   r6   r<   r;   r:   r   r   r   rw   r   r   r   r=   r   ro   r   rh   r   _inject_asgi_headersry   r   )r   r   r   r   r   r   r4   r4   r5   r     sH    


zLimiter._inject_asgi_headers)r?   rO   limitsrA   c                 C   sN  d }d }|D ]$}|j p|}|jr$q|jd ur@|j |jvr@q|jrT|d|j 7 }dt|jj	
 v rv||}n| }||g}	t|	r | jr| jg|	 }	|r|j|d k r|j|	f}t|jr||n|j}
| jj|jg|	R d|
is2| jd|j|| |}|j|	f} q4q| jd|j qq||j_|rJt|d S )Nz:%sr?   r   costz*ratelimit %s (%s) exceeded at endpoint: %sz4Skipping limit: %s. Empty value found in parameters.)scopeZ	is_exemptmethodsmethodlower
per_methodinspect	signaturerP   
parameterskeysallr{   r   callabler   rG   hitrh   r   rB   rF   rI   r   )r   r?   rO   r   Zfailed_limitZlimit_for_headerlimZlimit_scopeZ	limit_keyargsr   r4   r4   r5   Z__evaluate_limits  sJ    




zLimiter.__evaluate_limitsc              	   C   sz   zt |}W n ttfy&   d }Y n0 |d urBtt| S zt|}W n tyh   tdY n0 tt | S )NzQRetry-After Header does not meet RFC2616 - value is not of http-date or int type.)r   	TypeError
ValueErrorr   r   mktime	timetuple)r   Zretry_header_valueZretry_after_dateZretry_after_intr4   r4   r5   r     s    

zLimiter._determine_retry_time)r?   endpoint_funcin_middlewarerA   c                 C   sj  |d p
d}|}|r&|j  d|j nd}| jdkr8|n|}|rd| jrd|| jv sdtdd | jD rhdS g }g }	|s|| jv r| j| ng }g }	|| jv r| j| D ]R}
z|		t
|
| W q ty } z| jd|| W Y d}~qd}~0 0 qzg }| jrb| jrb|r"|| jv r"n@|  rR| j rR| jd	 d
| _d| _nt
tj| j }|s||	 }|rt
tj| j ng }||7 }tdd |D }|s|r|| jv r|r|t
tj| j 7 }| ||| W n~ tyd } zdt|t r | j!r8| js8| j"d d| _| #||| n| j$rN| j%d n W Y d}~n
d}~0 0 dS )z;
        Determine if the request is within limits
        rj   rM   .rN   c                 s   s   | ]}| V  qd S ra   r4   ).0fnr4   r4   r5   	<genexpr>=      z/Limiter._check_request_limit.<locals>.<genexpr>Nz2failed to load ratelimit for view function %s (%s)zRate limit storage recoveredFr   c                 s   s   | ]}|j  V  qd S ra   )override_defaults)r   r   r4   r4   r5   r   i  s   r   Tz&Failed to rate limit. Swallowing error)&r"   r!   r|   r]   rs   anyrt   r~   r   r}   listZwith_requestr   rh   rB   r   r   r   _Limiter__should_check_backendr   checkr   r   	itertoolschainro   rn   r   rm   _Limiter__evaluate_limitsr   r   r   rq   warn_check_request_limitry   r   )r   r?   r   r   Zendpoint_urlZ	view_funcZendpoint_func_nameZ_endpoint_keyr   Zdynamic_limitsr   eZ
all_limitsZroute_limitsZcombined_defaultsr4   r4   r5   r   &  s    	
zLimiter._check_request_limitr   )limit_valuerP   sharedr   r   r   error_messageexempt_whenr   r   rA   c                    s<   |r|nd  t dtf d 	f
dd}|S )N.)funcc                    s  pj } j d j }d }g }trFt|
		}nVz tt|
		}W n4 ty } zjd|| W Y d }~n
d }~0 0 j	
|g   |rʈj
|g | nj
|g | d }t }t|j D ]*\}|jdks|jdkr|j} q6qtd  dt rjt tttd fdd	}	|	S t tttd fd
d}
|
S d S )Nr   z*Failed to configure throttling for %s (%s)r?   Z	websocketz2No "request" or "websocket" argument on function "")r   kwargsrA   c                     s   j rZ|d| r|  nd }t|ts0tdjrZt|jddsZ| d d|j_	 | i |I d H }j rt|t
s|d|jj n||jj |S Nr?   zEparameter `request` must be an instance of starlette.requests.Request_rate_limiting_completeFTrJ   r]   r   r   r   r   rx   getattrrF   r   r   r   rH   rI   r   r   r?   rJ   r   idxr   r4   r5   async_wrapper  s*    

zCLimiter.__limit_decorator.<locals>.decorator.<locals>.async_wrapperc                     s   j rZ|d| r|  nd }t|ts0tdjrZt|jddsZ| d d|j_	 | i |}j rt|t
s|d|jj n||jj |S r   r   r   r   r4   r5   sync_wrapper  s*    

zBLimiter.__limit_decorator.<locals>.decorator.<locals>.sync_wrapper)rz   r"   r!   r   r   r   r   rh   rB   r   
setdefaultr   r   r~   r}   r   r   	enumerater   valuesnamer   asyncioiscoroutinefunction	functoolsr   r   r   )r   Zkeyfuncr   Zdynamic_limitZstatic_limitsr   Zconnection_typesigZ	parameterr   r   
Z_scoper   r   r   rP   r   r   r   r   r   )r   r   r5   	decorator  sp    


z,Limiter.__limit_decorator.<locals>.decorator)r   r   )r   r   rP   r   r   r   r   r   r   r   r   r   r4   r   r5   Z__limit_decorator  s    ,wzLimiter.__limit_decorator)	r   rP   r   r   r   r   r   r   rA   c	           	   
   C   s   | j ||||||||dS )a  
        Decorator to be used for rate limiting individual routes.

        * **limit_value**: rate limit string or a callable that returns a string.
         :ref:`ratelimit-string` for more details.
        * **key_func**: function/lambda to extract the unique identifier for
         the rate limit. defaults to remote address of the request.
        * **per_method**: whether the limit is sub categorized into the http
         method of the request.
        * **methods**: if specified, only the methods in this list will be rate
         limited (default: None).
        * **error_message**: string (or callable that returns one) to override the
         error message used in the response.
        * **exempt_when**: function returning a boolean indicating whether to exempt
        the route from the limit
        * **cost**: integer (or callable that returns one) which is the cost of a hit
        * **override_defaults**: whether to override the default limits (default: True)
        )r   r   r   r   r   r   _Limiter__limit_decorator)	r   r   rP   r   r   r   r   r   r   r4   r4   r5   r     s    zLimiter.limit)r   r   rP   r   r   r   r   rA   c              
   C   s   | j ||d|||||dS )ab  
        Decorator to be applied to multiple routes sharing the same rate limit.

        * **limit_value**: rate limit string or a callable that returns a string.
         :ref:`ratelimit-string` for more details.
        * **scope**: a string or callable that returns a string
         for defining the rate limiting scope.
        * **key_func**: function/lambda to extract the unique identifier for
         the rate limit. defaults to remote address of the request.
        * **per_method**: whether the limit is sub categorized into the http
         method of the request.
        * **methods**: if specified, only the methods in this list will be rate
         limited (default: None).
        * **error_message**: string (or callable that returns one) to override the
         error message used in the response.
        * **exempt_when**: function returning a boolean indicating whether to exempt
        the route from the limit
        * **cost**: integer (or callable that returns one) which is the cost of a hit
        * **override_defaults**: whether to override the default limits (default: True)
        T)r   r   r   r   r   )r   r   r   rP   r   r   r   r   r4   r4   r5   shared_limit7  s    zLimiter.shared_limitc                    sZ   d j  jf }| j| t r>t  fdd}|S t  fdd}|S dS )zF
        Decorator to mark a view as exempt from rate limits.
        z%s.%sc                     s    | i |I d H S ra   r4   akobjr4   r5   Z__async_innerj  s    z%Limiter.exempt.<locals>.__async_innerc                     s    | i |S ra   r4   r   r   r4   r5   Z__innerq  s    zLimiter.exempt.<locals>.__innerN)r"   r!   rs   addr   r   r   )r   r   r   Z_Limiter__async_innerZ_Limiter__innerr4   r   r5   exempt`  s    
zLimiter.exempt)N)T)	NFNFNNNr   T)NFNNNr   T)NNNr   T)'r!   r"   r#   __doc__r   r   r	   StrOrCallableStrboolr
   r   r   r   r   r   r   r   r   propertyr   rG   r   r   r   rH   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r4   r4   r4   r5   rL   Y   s   +

 H
	
65
0 d         

        
,     )rL   )9r   r   r   r   r   rf   ri   r   r   email.utilsr   r   r   typingr   r   r   r	   r
   r   r   r   r   r   r   Zlimits.errorsr   Zlimits.storager   r   Zlimits.strategiesr   r   Zstarlette.configr   Zstarlette.datastructuresr   Zstarlette.requestsr   Zstarlette.responsesr   r   Ztyping_extensionsr   errorsr   Zwrappersr   r   r   r   r   r    r6   r   rK   rL   r4   r4   r4   r5   <module>   s:   ,