o
    Df                     @   s   d dl Z d dlmZ G dd dedZG dd deedZG dd	 d	eedZG d
d dZG dd dZG dd deZ	G dd deZ
dS )    N)ABCMetac                   @   s4   e Zd ZdZdZdZdZdZdZdd Z	dd Z
dS )	JWEAlgorithmBasez+Base interface for all JWE algorithms.
    NJWEalgc                 C      t NNotImplementedError)selfraw_data r   T/home/ubuntu/webapp/venv/lib/python3.10/site-packages/authlib/jose/rfc7516/models.pyprepare_key      zJWEAlgorithmBase.prepare_keyc                 C   r   r   r   )r
   enc_algkeyr   r   r   generate_preset   r   z JWEAlgorithmBase.generate_preset)__name__
__module____qualname____doc__EXTRA_HEADERSnamedescriptionalgorithm_typealgorithm_locationr   r   r   r   r   r   r      s    r   )	metaclassc                   @   s"   e Zd ZdZdddZdd ZdS )JWEAlgorithmzInterface for JWE algorithm conforming to RFC7518.
    JWA specification (RFC7518) SHOULD implement the algorithms for JWE with this base implementation.
    Nc                 C   r   r   r   )r
   r   headersr   presetr   r   r   wrap   r   zJWEAlgorithm.wrapc                 C   r   r   r   )r
   r   ekr   r   r   r   r   unwrap   r   zJWEAlgorithm.unwrapr   )r   r   r   r   r    r"   r   r   r   r   r      s    
r   c                   @   s6   e Zd ZdZdddZdd ZdddZdd	d
ZdS )$JWEAlgorithmWithTagAwareKeyAgreementzInterface for JWE algorithm with tag-aware key agreement (in key agreement with key wrapping mode).
    ECDH-1PU is an example of such an algorithm.
    Nc                 C   r   r   r   )r
   r   r   
sender_keyr   r   r   r   !generate_keys_and_prepare_headers%   r   zFJWEAlgorithmWithTagAwareKeyAgreement.generate_keys_and_prepare_headersc                 C   r   r   r   )r
   r   r   r   r$   epkcektagr   r   r   agree_upon_key_and_wrap_cek(   r   z@JWEAlgorithmWithTagAwareKeyAgreement.agree_upon_key_and_wrap_cekc                 C   r   r   r   )r
   r   r   r   r$   r   r   r   r   r    +   r   z)JWEAlgorithmWithTagAwareKeyAgreement.wrapc                 C   r   r   r   )r
   r   r!   r   r   r$   r(   r   r   r   r"   .   r   z+JWEAlgorithmWithTagAwareKeyAgreement.unwrapr   )r   r   r   r   r%   r)   r    r"   r   r   r   r   r#   !   s    

r#   c                   @   sL   e Zd ZdZdZdZdZdZdZdd Z	dd Z
dd	 Zd
d Zdd ZdS )JWEEncAlgorithmNr   encc                 C      t | jd S N   )osurandomCEK_SIZEr
   r   r   r   generate_cek;      zJWEEncAlgorithm.generate_cekc                 C   r,   r-   )r/   r0   IV_SIZEr2   r   r   r   generate_iv>   r4   zJWEEncAlgorithm.generate_ivc                 C   s   t |d | jkrtdd S )Nr.   zInvalid "iv" size)lenr5   
ValueError)r
   ivr   r   r   check_ivA   s   zJWEEncAlgorithm.check_ivc                 C   r   )a  Encrypt the given "msg" text.

        :param msg: text to be encrypt in bytes
        :param aad: additional authenticated data in bytes
        :param iv: initialization vector in bytes
        :param key: encrypted key in bytes
        :return: (ciphertext, tag)
        r   )r
   msgaadr9   r   r   r   r   encryptE   s   	zJWEEncAlgorithm.encryptc                 C   r   )a8  Decrypt the given cipher text.

        :param ciphertext: ciphertext in bytes
        :param aad: additional authenticated data in bytes
        :param iv: initialization vector in bytes
        :param tag: authentication tag in bytes
        :param key: encrypted key in bytes
        :return: message
        r   )r
   
ciphertextr<   r9   r(   r   r   r   r   decryptP   s   
zJWEEncAlgorithm.decrypt)r   r   r   r   r   r   r   r5   r1   r3   r6   r:   r=   r?   r   r   r   r   r*   2   s    r*   c                   @   s,   e Zd ZdZdZdZdZdd Zdd ZdS )JWEZipAlgorithmNr   zipc                 C   r   r   r   r
   sr   r   r   compressc   r   zJWEZipAlgorithm.compressc                 C   r   r   r   rB   r   r   r   
decompressf   r   zJWEZipAlgorithm.decompress)	r   r   r   r   r   r   r   rD   rE   r   r   r   r   r@   ]   s    r@   c                       s4   e Zd ZdZ fddZdd Zedd Z  ZS )JWESharedHeaderziShared header object for JWE.

    Combines protected header and shared unprotected header together.
    c                    sR   i }|r	| | |r| | t | |r|ni | _|r$|| _d S i | _d S r   )updatesuper__init__	protectedunprotected)r
   rJ   rK   obj	__class__r   r   rI   o   s   

zJWESharedHeader.__init__c                 C   s   |  | | j | d S r   )rG   rJ   )r
   additionr   r   r   update_protectedy   s   
z JWESharedHeader.update_protectedc                 C   s$   t || r|S | |d|dS )NrJ   rK   )
isinstanceget)clsrL   r   r   r   	from_dict}   s   
zJWESharedHeader.from_dict)	r   r   r   r   rI   rP   classmethodrT   __classcell__r   r   rM   r   rF   j   s    
rF   c                       s    e Zd ZdZ fddZ  ZS )	JWEHeaderzHeader object for JWE.

    Combines protected header, shared unprotected header and specific recipient's unprotected header together.
    c                    sn   i }|r	| | |r| | |r| | t | |r!|ni | _|r(|ni | _|r2|| _d S i | _d S r   )rG   rH   rI   rJ   rK   header)r
   rJ   rK   rX   rL   rM   r   r   rI      s   


zJWEHeader.__init__)r   r   r   r   rI   rV   r   r   rM   r   rW      s    rW   )r/   abcr   r   r   r#   r*   r@   dictrF   rW   r   r   r   r   <module>   s    +