o
    ÙKÕiQ   ã                
   @   s|  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 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mZ d dlmZmZ d dlmZ d d	lmZ e ¡ Z d
Z!G dd„ dej"ƒZ#dd„ Z$de%fdd„Z&d1dd„Z'dd„ Z(dd„ Z)d1dd„Z*dd„ Z+dd„ Z,dd„ Z-d ed!dfd"d#„Z.d ed$ed%e/d&ee/ d!df
d'd(„Z0d)d*„ Z1d2d,ee/ fd-d.„Z2d/d0„ Z3dS )3é    N©Úwraps)ÚOptional)
ÚactionsÚapiÚdaemonÚentitlementsÚevent_loggerÚ
exceptionsÚlockÚmessagesÚstatusÚutil)ÚCVE_OR_USN_REGEX)Ú_is_attached)ÚAptProxyScopeÚsetup_apt_proxy)ÚUAConfig)Úmachine_tokenz[Kc                	   @   sJ   e Zd Zdefdd„Zdededee dee fdd	„Zd
efdd„Z	dS )ÚCLIEnableDisableProgressÚ
assume_yesc                C   s   | | _ || _d S ©N)Úis_interactiver   )Úselfr   © r   ú7/usr/lib/python3/dist-packages/uaclient/cli/cli_util.pyÚ__init__"   s   
z!CLIEnableDisableProgress.__init__Útotal_stepsÚ
done_stepsÚprevious_step_messageÚcurrent_step_messagec                C   s   |d ur
t |ƒ d S d S r   )Úprint)r   r   r   r   r    r   r   r   Úprogress&   s   ÿz!CLIEnableDisableProgress.progressÚeventc                 C   s:   |dkr
t |ƒ d S |dkrt || j¡st ¡ ‚d S d S )NÚinfoÚmessage_operation)r!   r   Úhandle_message_operationsr   r
   ÚPromptDeniedError)r   r#   Úpayloadr   r   r   Ú	_on_event1   s   þz"CLIEnableDisableProgress._on_eventN)
Ú__name__Ú
__module__Ú__qualname__Úboolr   Úintr   Ústrr"   r)   r   r   r   r   r   !   s    ýüû
úr   c                  O   s   d S r   r   ©ÚargsÚkwargsr   r   r   Ú_null_print9   s   r3   Újson_outputc                 C   s   | rt S tS r   )r3   r!   )r4   r   r   r   Ú&create_interactive_only_print_function=   s   r5   c                    ó   ‡ fdd„}|S )z1Decorator asserting exclusive access to lock filec                    ó   t ˆ ƒ‡ ‡fdd„ƒ}|S )Nc                    sF   t jˆdd�� ˆ |d| i|¤Ž}W d   ƒ |S 1 sw   Y  |S )Né   )Úlock_holderÚ
sleep_timeÚcfg)r   Ú	RetryLock)r;   r1   r2   Úretval)Úfr9   r   r   Únew_fH   s   
ÿþz0assert_lock_file.<locals>.wrapper.<locals>.new_fr   ©r>   r?   ©r9   ©r>   r   ÚwrapperG   s   z!assert_lock_file.<locals>.wrapperr   )r9   rC   r   rA   r   Úassert_lock_fileD   s   	rD   c                    ó   t ˆ ƒ‡ fdd„ƒ}|S )zDecorator asserting root userc                     s   t  ¡ st ¡ ‚ˆ | i |¤ŽS r   )r   Úwe_are_currently_rootr
   ÚNonRootUserErrorr0   rB   r   r   r?   V   s   zassert_root.<locals>.new_fr   r@   r   rB   r   Úassert_rootS   s   rH   c                    rE   )z>Decorator to verify if correct params are used for json formatc                    sH   | sˆ | g|¢R i |¤ŽS | j dkr| jst ¡ ‚ˆ | g|¢R i |¤ŽS )NÚjson)Úformatr   r
   ÚCLIJSONFormatRequireAssumeYes)Úcmd_argsr1   r2   rB   r   r   r?   c   s
   z&verify_json_format_args.<locals>.new_fr   r@   r   rB   r   Úverify_json_format_args`   s   	rM   c                    r6   )z“Decorator asserting attached config.
    :param msg_function: Optional function to generate a custom message
    if raising an UnattachedError
    c                    r7   )Nc                    sR   t |ƒjsˆrt| ddƒ}t| ddƒ}ˆ|||d� nt ¡ ‚ˆ | fd|i|¤ŽS )NÚcommandÚ Úservice)rN   Úservice_namesr;   r;   )r   Úis_attachedÚgetattrr
   ÚUnattachedError)r1   r;   r2   rN   rQ   )r>   Úraise_custom_error_functionr   r   r?   w   s   
ÿz/assert_attached.<locals>.wrapper.<locals>.new_fr   r@   ©rU   rB   r   rC   v   s   z assert_attached.<locals>.wrapperr   )rU   rC   r   rV   r   Úassert_attachedp   s   rW   c                    rE   )z&Decorator asserting unattached config.c                    s<   t |ƒjrt ¡ }tj|j dd¡d�‚ˆ | fd|i|¤ŽS )NÚnamerO   )Úaccount_namer;   )r   rR   r   Úget_machine_token_filer
   ÚAlreadyAttachedErrorÚaccountÚget)r1   r;   r2   Úmachine_token_filerB   r   r   r?   Œ   s   
ÿz"assert_not_attached.<locals>.new_fr   r@   r   rB   r   Úassert_not_attached‰   s   r_   c                    r6   )Nc                    s   t ˆ ƒ‡‡ fdd„ƒ}|S )Nc                    s:   t | ddƒ}t t|¡stj|ˆ d�‚ˆ| fd|i|¤ŽS )NÚsecurity_issuerO   )ÚissueÚcmdr;   )rS   ÚreÚmatchr   r
   ÚInvalidSecurityIssueIdFormat)r1   r;   r2   r`   )rb   r>   r   r   r?   š   s   þz@assert_vulnerability_issue_valid.<locals>.wrapper.<locals>.new_fr   r@   ©rb   rB   r   rC   ™   s   
z1assert_vulnerability_issue_valid.<locals>.wrapperr   )rb   rC   r   rf   r   Ú assert_vulnerability_issue_valid˜   s   rg   c                 C   sf   t j||d�\}}|r|rtjd |¡| d |¡dd�‚|r(tjd |¡| d�‚tj| d |¡dd�‚)z£Raises a custom error for enable/disable commands when unattached.

    Takes into consideration if the services exist or not, and notify the user
    accordingly.)Únamesr;   z, rO   )Úvalid_serviceÚ	operationÚinvalid_serviceÚservice_msg)ri   rj   )rj   rk   rl   )r   Úget_valid_entitlement_namesr
   ÚUnattachedMixedServicesErrorÚjoinÚUnattachedValidServicesErrorÚUnattachedInvalidServicesError)rN   rQ   r;   Úentitlements_foundÚentitlements_not_foundr   r   r   Ú&_raise_enable_disable_unattached_errorª   s(   ýü
ÿýrt   r;   Úreturnc                 C   s|   t  | ¡}|j}|rt tjj|d�¡ nt tj¡ t	 
¡  t	 | ¡ t | ¡\}}t |¡}t t |¡¡ t ¡  d S )N)Úcontract_name)r   rZ   rv   r#   r$   r   ÚATTACH_SUCCESS_TMPLrJ   ÚATTACH_SUCCESS_NO_CONTRACT_NAMEr   ÚstopÚcleanupr   r   Úformat_tabularr   Úhandle_unicode_charactersÚprocess_events)r;   r^   rv   Ústatus_dictÚ_retÚoutputr   r   r   Úpost_cli_attachÆ   s   
ÿ

r�   ÚscopeÚset_keyÚ	set_valuec                 C   sR   |t jkr| j}| j}n|t jkr| j}| j}d|v r|}n|}t|||d� dS )zS
    Handles setting part the apt proxies - global and uaclient scoped proxies
    Úhttps)Ú
http_proxyÚhttps_proxyÚproxy_scopeN)r   ÚGLOBALÚglobal_apt_http_proxyÚglobal_apt_https_proxyÚUACLIENTÚua_apt_http_proxyÚua_apt_https_proxyr   )r;   r‚   rƒ   r„   r†   r‡   r   r   r   Úconfigure_apt_proxyÚ   s   
	

ÿr�   c                 C   s`   t  g d¢¡}|  ¡ s.tj d |t|ƒt¡¡ tj 	¡  t
 d¡ tj d¡ |  ¡ rd S d S )N)ú|ú/ú-ú\z{}{}{}gš™™™™™¹?ú)Ú	itertoolsÚcycleÚis_setÚsysÚstdoutÚwriterJ   ÚnextÚCLEAR_LINE_ANSI_CODEÚflushÚtimeÚsleep)Ústop_spinnerÚmsgÚspinnerr   r   r   Úrun_spinnerò   s   ÿ

úr£   rO   r¡   c                    r6   )Nc                    r7   )Nc                   s‚   t j ¡ sˆ | fd|i|¤ŽS t ¡ }tjt|ˆfd�}| ¡  ˆ | fd|i|¤Ž}| ¡  | 	¡  t j 
dt ¡ t j ¡  |S )Nr;   )Útargetr1   ú)r˜   r™   ÚisattyÚ	threadingÚEventÚThreadr£   ÚstartÚsetro   rš   rœ   r�   )r1   r;   r2   r    Úspinner_threadr=   )r>   r¡   r   r   r?   ÿ   s   
ÿ
z,with_spinner.<locals>.wrapper.<locals>.new_fr   r@   ©r¡   rB   r   rC   þ   s   zwith_spinner.<locals>.wrapperr   )r¡   rC   r   r­   r   Úwith_spinnerý   s   r®   c                 C   st   | dkrt jj|  t jj S | dkrt jj|  t jj S | dkr*t jj|  t jj S | dkr8t jj|  t jj S | S )NÚlowÚmediumÚhighÚcritical)r   ÚTxtColorÚINFOBLUEÚENDCÚWARNINGYELLOWÚORANGEÚFAIL)Úpriorityr   r   r   Úcolorize_priority  s   ÿrº   r   )rO   )4r•   rc   r˜   r§   rž   Ú	functoolsr   Útypingr   Úuaclientr   r   r   r   r	   r
   r   r   r   r   Ú'uaclient.api.u.pro.security.fix._commonr   Ú(uaclient.api.u.pro.status.is_attached.v1r   Úuaclient.aptr   r   Úuaclient.configr   Úuaclient.filesr   Úget_event_loggerr#   rœ   ÚAbstractProgressr   r3   r-   r5   rD   rH   rM   rW   r_   rg   rt   r�   r/   r�   r£   r®   rº   r   r   r   r   Ú<module>   sP    0

ÿþýü
û