U
    Jhn                     @   s   d dl Z d dlmZmZmZmZmZ d dlmZm	Z	m
Z
mZmZ d dlmZ d dlmZmZ d dlmZmZmZ e	 Ze eeZddiZG d	d
 d
ejZG dd deZG dd deZ G dd deZ!G dd deZ"G dd deZ#dS )    N)AnyDictOptionalTupleType)aptevent_loggermessagessystemutil)repo)EntitlementWithMessageUAEntitlement)MessagingOperationsMessagingOperationsDictStaticAffordanceZGenuineIntelZintelc                   @   s   e Zd ZdZejZejZej	Z
ejjZdZdZdZdZdZedddZeeee f dd	d
Zedd Zeeedf dddZeeedf dddZ ee!dddZ"ddddZ#dS )RealtimeKernelEntitlementzrealtime-kernelzubuntu-pro-realtime-kernel.gpgTFZUbuntuRealtimeKernelreturnc                 C   s(   t jt| jtdgd}t| |S )z%Check if system needs to be rebooted.zlinux-.*-realtime)Zinstalled_pkgsZinstalled_pkgs_regex)r
   Zshould_rebootsetpackageseventZneeds_reboot)selfZreboot_required r   @/usr/lib/python3/dist-packages/uaclient/entitlements/realtime.py_check_for_reboot   s    
z+RealtimeKernelEntitlement._check_for_rebootc                 C   s   t jt tjttjttjtiS N)GenericRealtimevariant_nameNvidiaTegraRealtimeRaspberryPiRealtimeIntelIotgRealtimer   r   r   r   _get_variants(   s        z'RealtimeKernelEntitlement._get_variantsc                 C   s   | j r
d S tS r   )access_onlyr   r"   r   r   r   default_variant0   s    z)RealtimeKernelEntitlement.default_variant.c                 C   s>   ddl m}m} ddlm} t|tjt|tjt|tj	fS )Nr   )FIPSEntitlementFIPSUpdatesEntitlement)LivepatchEntitlement)
Zuaclient.entitlements.fipsr&   r'   Zuaclient.entitlements.livepatchr(   r   r	   ZREALTIME_FIPS_INCOMPATIBLEZ"REALTIME_FIPS_UPDATES_INCOMPATIBLEZREALTIME_LIVEPATCH_INCOMPATIBLE)r   r&   r'   r(   r   r   r   incompatible_services6   s      z/RealtimeKernelEntitlement.incompatible_servicesc                 C   s   t jdd dffS )Nc                   S   s   t  S r   )r
   Zis_containerr   r   r   r   <lambda>P       z>RealtimeKernelEntitlement.static_affordances.<locals>.<lambda>F)r	   Z#REALTIME_ERROR_INSTALL_ON_CONTAINERr"   r   r   r   static_affordancesK   s
    z,RealtimeKernelEntitlement.static_affordancesc                 C   sD   d }| j stjtjddfg}d }| js:tjdtjifg}||dS )NT)msgdefaultr-   )
pre_enablepre_disable)r$   r   prompt_for_confirmationr	   ZREALTIME_PROMPTpurgeZREALTIME_PRE_DISABLE_PROMPT)r   r/   r0   r   r   r   	messagingU   s&    
 
z#RealtimeKernelEntitlement.messagingNc                 C   s<   t | jt t }|r8tt|tjj	| j
d d S )N)title)r   r   intersectionr   Zget_installed_packages_namesremove_packageslistr	   ZDISABLE_FAILED_TMPLformatr4   )r   r   r   r   r   r6   u   s    

z)RealtimeKernelEntitlement.remove_packages)$__name__
__module____qualname__namer	   ZREALTIME_TITLEr4   ZREALTIME_DESCRIPTIONdescriptionZREALTIME_HELP_TEXTZ	help_textZurlsZREALTIME_HOME_PAGEZhelp_doc_urlZrepo_key_fileZapt_noninteractiveZsupports_access_onlyZsupports_purgeoriginboolr   r   strr   r   r#   propertyr%   r   r   r)   r   r,   r   r3   r6   r   r   r   r   r      s,   	
	r   c                       s.   e Zd Zeeedf d fddZ  ZS )RealtimeVariant.r   c                    s(   t  fdd j D }t j| S )Nc                    s*   g | ]"\}}t |tjj j|jd qS ))servicevariant)r   r	   ZREALTIME_VARIANT_INCOMPATIBLEr8   r4   ).0r<   clsr"   r   r   
<listcomp>   s    z9RealtimeVariant.incompatible_services.<locals>.<listcomp>)tupleZother_variantsitemssuperr)   )r   Zincompatible_variants	__class__r"   r   r)      s    
z%RealtimeVariant.incompatible_services)r9   r:   r;   rA   r   r   r)   __classcell__r   r   rK   r   rB      s   rB   c                       s>   e Zd ZdZejZejZdZ	dZ
eed fddZ  ZS )r   genericTr   c                    sd   t  j}t j}|dkr`|dp&g }tjj| j	| j
d|p>dd}|tjd|if ||d< |S )NrN   r/   unknown)rD   rC   Zbase_flavorcurrent_flavorr-   )rJ   r3   r
   Zget_kernel_infoZflavorgetr	   Z#KERNEL_FLAVOR_CHANGE_WARNING_PROMPTr8   r   r<   appendr   r1   )r   r3   rP   r/   r-   rK   r   r   r3      s"    
zGenericRealtime.messaging)r9   r:   r;   r   r	   ZREALTIME_GENERIC_TITLEr4   ZREALTIME_GENERIC_DESCRIPTIONr=   
is_variantcheck_packages_are_installedrA   r   r3   rM   r   r   rK   r   r      s   r   c                   @   s$   e Zd ZdZejZejZdZ	dZ
dS )r   znvidia-tegraTN)r9   r:   r;   r   r	   ZREALTIME_NVIDIA_TITLEr4   ZREALTIME_NVIDIA_DESCRIPTIONr=   rS   rT   r   r   r   r   r      s
   r   c                   @   s2   e Zd ZdZejZejZdZ	dZ
edddZdS )r    ZraspiTr   c              
   C   sb   d}z$t |  }d|kp&d|kW S  tk
r\ } ztd| W Y dS d }~X Y nX d S )Nz/proc/device-tree/modelzraspberry pi 4zraspberry pi 5z)Error while detecting if raspberry pi: %rF)r
   Z	load_filestriplower	ExceptionLOGinfo)r   Zproc_file_pathZmodeler   r   r   variant_auto_select   s    z'RaspberryPiRealtime.variant_auto_selectN)r9   r:   r;   r   r	   ZREALTIME_RASPI_TITLEr4   ZREALTIME_RASPI_DESCRIPTIONr=   rS   rT   r?   r[   r   r   r   r   r       s   r    c                   @   sJ   e Zd ZdZejZejZdZ	dZ
eeef eeeej f dddZdS )r!   z
intel-iotgT)platform_checksr   c                 C   sT   t  jpd}t||}|dg }||kr2dS dtjj| j|d	|dfS d S )N cpu_vendor_ids)TNF,)r4   ZvendorZsupported_vendors)
r
   Zget_cpu_infoZcpuinfo_vendor_idCPU_VENDOR_MAPrQ   r	   ZINAPPLICABLE_VENDOR_NAMEr8   r4   join)r   r\   Z	vendor_idr^   r   r   r   verify_platform_checks   s    z(IntelIotgRealtime.verify_platform_checksN)r9   r:   r;   r   r	   ZREALTIME_INTEL_TITLEr4   ZREALTIME_INTEL_DESCRIPTIONr=   rS   rT   r   r@   r   r   r?   r   ZNamedMessagerb   r   r   r   r   r!      s   
r!   )$Zloggingtypingr   r   r   r   r   Zuaclientr   r   r	   r
   r   Zuaclient.entitlementsr   Zuaclient.entitlements.baser   r   Zuaclient.typesr   r   r   Zget_event_loggerr   Z	getLoggerZreplace_top_level_logger_namer9   rX   r`   ZRepoEntitlementr   rB   r   r   r    r!   r   r   r   r   <module>   s   m