U
    ^gU\                     @   s  d dl 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 dlmZ d dlZd dl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 d d
lmZ d dl m!Z! d dl"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/m0Z0 d dl1m2Z2m3Z3m4Z4 d dl5m6Z6 d dl7m8Z8 e%ddddddddddddZ9eee:ef dddZ;e0e'd d!d"Z<ee:d#d$d%Z=e0e:e8e:e:e:d&d'd(Z>dXe8e?e:d*d+d,Z@ddde3jAdd-d.fee0ee	e8 d/f e
e: e
e: e
e. e3e
e: e:e
e: e'd0d1d2ZBdYeee
e: e
e: e
e: e
e: d3d4d5ZCejd6d7dd8d9d:ZDeDjEd;d7edd8d<d=ZFeDjEd>d7ejGd?dd@dAejGdBddCdAejGdDddEdAejGdFddGdAejGdHdIdJdAedZe
e: e
e: e:e
e: e?ddKdLdMZHeDjEdNd7ejGd?dd@dAejGdBddCdAejGdDddEdAejGdFdd)dOdPejGdQd)dRdSed[e:e
e: e
e: e:e?ddTdUdVZIeJdWkreD  dS )\    N)partial)
AnyCallableDict	GeneratorListOptionalSetTupleUnioncast)urlparse)Flaskcurrent_app)with_appcontext)RateLimiter)Consolegroup)Live)Pretty)Confirm)Table)Theme)Tree)	TypedDict)MethodNotAllowedNotFound)Rule)Limiter)
ConfigVarsExemptionScopeHeaderNames)get_qualified_name)Limitz
bold greenbold redZmagentaZcyanZyellowzbold yellow)successZdangererror	blueprintdefaultcallableentityZexemptroutehttpoption)funcreturnc                 C   s>   t | r6| jdkr d| j dS d| j d| j dS t| S )Nz<lambda>z[callable]<lambda>(z)[/callable]z
[callable].z()[/callable])r)   __name__
__module__r   )r.    r3   :/tmp/pip-unpacked-wheel-8ut481ky/flask_limiter/commands.pyrender_func7   s
    
r5   )extr/   c                 C   s~   t | jp
d}| jrz|d| jjj d |d| jj d |t| jpPi  | j }|rp|d n
|d |S )NzN/A[entity]	[/entity]z[success]OK[/success]z[error]Error[/error])	r   Z_storage_uristorageadd	__class__r1   r   Z_storage_optionscheck)r6   renderZhealthr3   r3   r4   render_storage?   s    

r>   )strategyr/   c                 C   s   d| j j dS )Nr7   r8   )r;   r1   )r?   r3   r3   r4   render_strategyM   s    r@   )limiterendpointlimitkeymethodr/   c                 C   s   || ||g}| jr&| jr*| j s*dS | jj|jf| }| jj|jf| }|spd|d  d|jj dS d|d  d|jj dS d S )Nz&: [error]Storage not available[/error]z: [error]Fail[/error] (   z out of z remaining)z: [success]Pass[/success] ()	scope_forr9   r<   rA   testrC   Zget_window_statsamount)rA   rB   rC   rD   rE   argsrH   statsr3   r3   r4   render_limit_stateQ   s    rL   T)rC   simpler/   c                 C   sl   t | j}|r|S g }| jr2|dt| j  | jrN|dt| j  |rh| dd| d}|S )Nzdeduct_when: zexempt_when: z
 [option]{z, z
}[/option])strrC   Zdeduct_whenappendr5   Zexempt_whenjoin)rC   rM   r=   optionsr3   r3   r4   render_limita   s    
rR   GET .)apprA   limitsrB   r'   ruleexemption_scoperH   rE   labelr/   c
                    s  d }
|r|r| d| }
|
p$|	p$d}	t |	 g }|d |d  D ]}|r| j|d }|rt||j| |krtdn"||j|rt|ndkrdnd}nd}|jr|r|jr|jD ]T}t	|d}d	| d
| d| d| d
	}|r|t
||pd|||7 }|| qqBt	|d}d	| d
| d| d
}|rT|t
||pHd|||7 }|| qB|sx|rx d n fdd|D   S )N: rT   r   rF   r'   r+   r(   F[]z [http](z
)[/http][/z[/z[exempt]Exempt[/exempt]c                    s   g | ]}  |qS r3   r:   ).0entry
renderabler3   r4   
<listcomp>   s     z!render_limits.<locals>.<listcomp>)r   view_functionsgetlimit_managerZblueprint_limitsZdecorated_limitsr"   
per_methodmethodsrR   rL   rO   r:   )rU   rA   rV   rB   r'   rW   rX   rH   rE   rY   _labelentriesrC   Z	view_funcsourcerenderedr_   r3   r`   r4   render_limitso   sd    

    
    rl   )rU   consolerB   pathrE   r/   c           	   
   C   s   |s|sd S |r2|t jkr|S |d| d n|r| jd}t|}z$|j|j||jd\}}t	t
|W S  tk
r   |d| d Y n8 tk
r   |st|d|  d| d Y nX td S )Nz[red]Error: z
 not foundzdev.null)rE   Z
query_argsz[error]Error: z, could not be matched to an endpoint[/error]rZ   )r   rc   printurl_mapbindr   matchrn   queryr   rN   r   r   AssertionErrorupper
SystemExit)	rU   rm   rB   rn   rE   adapterparsedfilter_endpoint_r3   r3   r4   get_filtered_endpoint   s4    
  


r{   z-Flask-Limiter maintenance & utility commmands)helpr/   c                   C   s   d S Nr3   r3   r3   r3   r4   cli   s    r   z View the extension configurationc            	   	   C   s  t   ttd} tt jdt }|o8t|d }|rttdd}|	d |	d |	d |
d	tjt|j |
d
tjt|j |
dtjt|j ttj}tt|j}|tj}|d |d |tj |d |t| |
d|| |jjr|
dtjtdd |jjD  |
d tjt|j  |
d tj!t|j" |
d tj#t|j$ |
d tj%t|j& n|
dtjtg  |jj'r4|
dtj(tdd |jj'D  |
d tj)t|j* |
d tj+t|j, |
d tj-t|j. |
d tj/t|j0 n|
dtj(tg  |j1rv|
dtj2tdd t3j4|j1 D  |j5r(ttj6}|tj7 |tj8 |tj9 |tj: tt|j5}|t|j;t<j=  |t|j;t<j>  |t|j;t<j?  |t|j@ |
d|| n|
dtj6td |
dtjAt|jB |
dtjCt|jD | E| n| jEdt  dd W 5 Q R X d S )NZthemerA   r   zFlask-Limiter Config)titleZNotesConfigurationValueZEnabledzKey Functionz
Key PrefixZInstanceBackendZStatuszRate Limiting ConfigApplication Limitsc                 S   s   g | ]}t |qS r3   rR   r^   rC   r3   r3   r4   rb      s   zconfig.<locals>.<listcomp>zApplicationLimits LimitszDefault Limitsc                 S   s   g | ]}t |qS r3   r   r   r3   r3   r4   rb     s   zMeta Limitsc                 S   s   g | ]}t |qS r3   r   r   r3   r3   r4   rb   C  s   zHeader configurationFzFail on first breachzOn breach callback(No Flask-Limiter extension installed on r$   style)Fr   test_request_contextr   limiter_themelist
extensionsrd   setr   Z
add_columnZadd_rowr   ZENABLEDr   ZenabledZKEY_FUNCr5   Z	_key_funcZ
KEY_PREFIXZ_key_prefixr   ZSTRATEGYr@   rA   r:   ZSTORAGE_URIZSTORAGE_OPTIONSr>   re   application_limitsZAPPLICATION_LIMITSZAPPLICATION_LIMITS_PER_METHODZ_application_limits_per_methodZAPPLICATION_LIMITS_EXEMPT_WHENZ_application_limits_exempt_whenZAPPLICATION_LIMITS_DEDUCT_WHENZ_application_limits_deduct_whenZAPPLICATION_LIMITS_COSTZ_application_limits_costZdefault_limitsZDEFAULT_LIMITSZDEFAULT_LIMITS_PER_METHODZ_default_limits_per_methodZDEFAULT_LIMITS_EXEMPT_WHENZ_default_limits_exempt_whenZDEFAULT_LIMITS_DEDUCT_WHENZ_default_limits_deduct_whenZDEFAULT_LIMITS_COSTZ_default_limits_cost_meta_limitsZMETA_LIMITS	itertoolschainZ_headers_enabledZHEADERS_ENABLEDZHEADER_RESETZHEADER_REMAININGZHEADER_RETRY_AFTERZHEADER_RETRY_AFTER_VALUEZ_header_mappingr!   RESETZ	REMAININGZRETRY_AFTERZ_retry_afterZFAIL_ON_FIRST_BREACHZ_fail_on_first_breachZ	ON_BREACHZ
_on_breachro   )	rm   limitersrA   Zextension_detailsZlimiter_configZlimiter_config_valuesnodeZheader_configsheader_valuesr3   r3   r4   config   s:   




      



  



  


  r   z3Enumerate details about all routes with rate limitsz
--endpointzEndpoint to filter by)r(   r|   z--pathzPath to filter byz--methodzHTTP Method to filter byz--keyzTest the limitz--watch/--no-watchFzCreate a live dashboard)rB   rn   rE   rD   watchr/   c                    s  t   t jdt }|r,t|d nd ttd}rj}i t	t | }t
t j|dd dD ]}	|	j}
|
dkrqrt|
dd	krd|
dd d
 }|g ttt |t |
||
||	|t |
|d
 qrdg ttt |t |
d|
d |	|t |
d d
 qrt td fdd}|sf||  nZt| |dddB}z||  td W n tk
r   Y qY nX qzW 5 Q R X n|jdt  dd W 5 Q R X d S )NrA   r   r   c                 S   s   t | S r~   rN   rr3   r3   r4   <lambda>      zlimits.<locals>.<lambda>rD   staticr0   rF   )rX   rE   rH   rootrT   r}   c                  3   s   rXj jrXsXsXttttjj g fddV  ttj jg fddV  D ]L} | dkr|tdtj	 d ntd|  d  fd	d
|  D   V  q\d S )Nz[gold3]Meta Limits[/gold3])rH   rE   rY   z![gold3]Application Limits[/gold3]r   z[gold3]z[/gold3]z[blue]z[/blue]c                    s   g | ]}  | qS r3   r]   )r^   ra   Z
group_treer3   r4   rb     s     z6limits.<locals>.console_renderable.<locals>.<listcomp>)
re   r   rl   r   r   r   r   r   r   name)r   rB   groupsrD   rA   rE   rn   r   r4   console_renderable  s<    
z"limits.<locals>.console_renderableg?T)rm   Zrefresh_per_secondZscreenr   r$   r   )r   r   r   rd   r   r   r   r   re   r{   sortedrp   
iter_rulesrB   lensplitrP   
setdefaultrO   r   rl   resolve_limitsrX   r   r   ro   r   updatetimesleepKeyboardInterrupt)rB   rn   rE   rD   r   r   rm   managerry   rW   rule_endpointbp_fullnamer   Zliver3   r   r4   rV   v  s    
    
 
      rV   zClear limits for a specific keyzKey to reset the limits for)r(   requiredr|   z-yzSkip prompt for confirmation)Zis_flagr|   )rD   rB   rn   rE   yr/   c                 C   s  t   tt jdt }|r,|d nd }ttd}|r|j}t	t ||||}	t
dtttt df d}
i }tt j|	dd d	D ]t}|j}|d
krqt|ddkrd|dd d }|
||t ||d||< q|
||t |dd||< qd }|	s|jj}|s~|r>|tt ||g fd| d | D ]6\}}|d rF|tt ||d ||d | d qF|std|  dr|rtd}|D ]4}|j|j | |!d| |"t#| d q|| | D ]\}}|d rt|}|d \}}|| D ]}|j$rv|d rv|d j%rv|sv|d j%D ] }|j|j | |!|| qRn|j|j | |!|| |"t#| d q || qn|jdt  dd W 5 Q R X d S )NrA   r   r   Details.)rW   rV   c                 S   s   t | S r~   r   r   r3   r3   r4   r     r   zclear.<locals>.<lambda>r   r   r0   rF   r   rT   r   )rY   rH   rV   rW   )rW   rH   z/Proceed with resetting limits for key: [danger]z
[/danger]?z: [success]Cleared[/success]r   r$   r   )&r   r   r   r   rd   r   r   r   re   r{   r   r   r
   r   r#   r   rp   r   rB   r   r   rP   r   r   ro   rl   itemsr   askr   rA   clearrC   rG   r:   rR   rf   rg   )rD   rB   rn   rE   r   r   rA   rm   r   ry   r   Zrule_limitsrW   r   r   r   detailsr   rC   r(   Z	decoratedZrule_methodr3   r3   r4   r     s    
    
 
  	









r   __main__)T)N)NNrS   NF)NNrS   F)Kr   r   	functoolsr   typingr   r   r   r   r   r   r	   r
   r   r   urllib.parser   ZclickZflaskr   r   Z	flask.clir   Zlimits.strategiesr   Zrich.consoler   r   Z	rich.liver   Zrich.prettyr   Zrich.promptr   Z
rich.tabler   Z
rich.themer   Z	rich.treer   Ztyping_extensionsr   Zwerkzeug.exceptionsr   r   Zwerkzeug.routingr   Zflask_limiterr   Zflask_limiter.constantsr   r    r!   Zflask_limiter.utilr"   Zflask_limiter.wrappersr#   r   rN   r5   r>   r@   rL   boolrR   NONErl   r{   r   commandr   r-   rV   r   r1   r3   r3   r3   r4   <module>   s   0    C #

  
     s
    u
