U
    Jh                     @   s   d dl Z d dlZd dlmZmZ d dlmZmZmZm	Z	m
Z
 d dlmZ dZdZeefZdZdZd	Zd
ZdZdZdZdZed Zee
eZG dd deZdS )    N)AnyDict)
exceptionshttpsecret_managersystemutil)PublicCloudAutoAttachInstancez169.254.169.254z[fd00:ec2::254]zhttp://{}/latest/api/tokenz0http://{}/latest/dynamic/instance-identity/pkcs7z3http://{}/latest/dynamic/instance-identity/documentz/sys/hypervisor/uuidz /sys/class/dmi/id/product_serialz/sys/class/dmi/id/product_uuidZ21600zX-aws-ec2-metadata-tokenz-ttl-secondsc                   @   s   e Zd ZdZdZedddZeej	e
jdddgdeeef dd	d
Zej	e
jdddgdedddZeedddZdd Zej	e
jdddgddd ZeedddZeedddZedddZeedddZdS )AWSAutoAttachInstanceN)urlc                 C   s6   t j||dd}|jdkr |jS tj|j|jdd S )N   )headerstimeout   codebody)r   readurlr   r   r   CloudMetadataError)selfr   r   response r   5/usr/lib/python3/dist-packages/uaclient/clouds/aws.py_get_imds_url_response   s    
 z,AWSAutoAttachInstance._get_imds_url_responseg      ?r   )Zretry_sleeps)returnc                 C   s6   |   }t| j}| j||d}tj| d|iS )Nr   Zpkcs7)_request_imds_v2_token_headersIMDS_URLformat_ip_addressr   r   secrets
add_secret)r   r   r   Zimds_url_responser   r   r   identity_doc*   s
    z"AWSAutoAttachInstance.identity_docc              
   C   sj   |   }t| j}zt| j||d}W n8 tjk
rd } zt	d| i  W Y S d}~X Y nX |S )a  
        Get the instance identity doc associated with the current instance.

        See
        https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/retrieve-iid.html
        for more context.

        @return: Dict containing the instance identity document.
        r   z-Error decoding instance identity document: %sN)
r   _IMDS_IID_URLr   r   jsonloadsr   ZJSONDecodeErrorLOGdebug)r   r   r   ii_docer   r   r   _get_ii_doc3   s    z!AWSAutoAttachInstance._get_ii_docc                 C   s0   |   }|dd}|dd}t|p.t|S )a  
        Determines if the instance is likely Ubuntu Pro.

        Criteria: if any billing-product or marketplace-product-code is
        present, then is likely a Pro instance.

        @return: Boolean indicating if the instance is likely pro or not.
        ZbillingProductsNZmarketplaceProductCodes)r*   getbool)r   r(   Zbilling_productsZmarketplace_product_codesr   r   r   is_likely_proI   s    
z#AWSAutoAttachInstance.is_likely_proc                 C   s~   t D ]X}z| j|d}W n8 tk
rP } ztd|t|dd W 5 d }~X Y qX || _ q^q| jd krztjd	t d|S )N)
ip_addressz*Could not reach AWS IMDS at http://%s: %s
reason z, )Z	addresses)
IMDS_IP_ADDRESS_get_imds_v2_token_headers	Exceptionr&   Zwarninggetattrr   r   ZAWSNoValidIMDSjoin)r   Zaddressr   r)   r   r   r   r   X   s     

z4AWSAutoAttachInstance._request_imds_v2_token_headers      c                 C   s   | j dkrd S | j rt| j iS tjt|dttidd}|jdkrd|j	| _ t
j| j  t| j iS |jdkrxd| _ d S tj|j|j	dd S )NZIMDSv1ZPUTr   )methodr   r   r   i  r   )
_api_tokenAWS_TOKEN_PUT_HEADERr   r   IMDS_V2_TOKEN_URLr   AWS_TOKEN_REQ_HEADERAWS_TOKEN_TTL_SECONDSr   r   r   r    r!   r   r   )r   r.   r   r   r   r   r2   k   s*    




 z0AWSAutoAttachInstance._get_imds_v2_token_headersc                 C   s   dS )NZawsr   r   r   r   r   
cloud_type   s    z AWSAutoAttachInstance.cloud_typec                 C   s   z$t t}d|dd kr"W dS W n tk
r8   Y nX t t }t t }d|dd   krz|dd krn ndS dS )z$This machine is a viable AWSInstanceZec2r      TF)r   Z	load_fileSYS_HYPERVISOR_PRODUCT_UUIDFileNotFoundErrorDMI_PRODUCT_UUIDlowerDMI_PRODUCT_SERIAL)r   Zhypervisor_uuidZdmi_uuidZ
dmi_serialr   r   r   	is_viable   s    

(zAWSAutoAttachInstance.is_viablec                 C   s   dS )ZUnsupportedFr   r>   r   r   r   should_poll_for_pro_license   s    z1AWSAutoAttachInstance.should_poll_for_pro_license)wait_for_changer   c                C   s   t  d S )N)r   ZInPlaceUpgradeNotSupportedError)r   rH   r   r   r   is_pro_license_present   s    z,AWSAutoAttachInstance.is_pro_license_present)__name__
__module____qualname__r9   r   strr   propertyr   Zretryr   r   r   r   r"   r*   r,   r-   r   r2   r?   rF   rG   rI   r   r   r   r   r
      s&   
r
   )r$   Zloggingtypingr   r   Zuaclientr   r   r   r   r   Zuaclient.cloudsr	   ZIMDS_IPV4_ADDRESSZIMDS_IPV6_ADDRESSr1   r;   r   r#   rA   rE   rC   r=   r:   r<   Z	getLoggerZreplace_top_level_logger_namerJ   r&   r
   r   r   r   r   <module>   s$   