o
    c                     @   s  d dl Z d dlmZ d dlmZ d dlmZ d dlmZ de	de	de j
e	 d	e	fd
dZ	dde	de	de jd	e	fddZde	de	de j
e	 d	e je	e j
e	 f fddZ	dde	de	de jd	e	fddZ	dde	de	de jd	e	fddZ	dde	de	de jd	e	fddZG dd deZdS )    N)Cipher)AES)ECB)bytes_eqwrapping_keyarreturnc                 C   s   t t| t  }t|}tdD ]2}t|D ]+}||||  }tj|d d dd|| | d A j	ddd}|dd  ||< qq|
 dksMJ |d| S )	N      big	byteorder   lengthr       )r   r   r   	encryptorlenrangeupdateint
from_bytesto_bytesfinalizejoin)r   r   r   r   njib r!   Jusr/lib/python3.10/site-packages/cryptography/hazmat/primitives/keywrap.py
_wrap_core   s   $
r#   key_to_wrapbackendc                    sn   t | dvr
tdt  dk rtdt  d dkr tdd} fd	d
tdt  dD }t| ||S )N          /The wrapping key must be a valid AES key lengthr'   z)The key to wrap must be at least 16 bytesr   r   z-The key to wrap must be a multiple of 8 bytes   c                       g | ]
} ||d   qS r   r!   .0r   r$   r!   r"   
<listcomp>5       z aes_key_wrap.<locals>.<listcomp>)r   
ValueErrorr   r#   )r   r$   r%   r   r   r!   r0   r"   aes_key_wrap&   s   r4   c           	      C   s   t t| t  }t|}ttdD ]6}tt|D ]-}tj|dd|| | d A j	ddd||  }|
|}|d d }|dd  ||< qq| dksSJ ||fS )	Nr
   r   r   r   r   r   r   r   )r   r   r   	decryptorr   reversedr   r   r   r   r   r   )	r   r   r   r5   r   r   r   Zatrr    r!   r!   r"   _unwrap_core9   s   

r7   c                    s   t | dvr
tddt  jddd }dt  d  d } d|   t  dkrFtt| t  }||  }| d	ksDJ |S  fd
dt	dt  dD }t
| ||S )Nr&   r*      YY   r   r   r       r   c                    r,   r-   r!   r.   r0   r!   r"   r1   e   r2   z-aes_key_wrap_with_padding.<locals>.<listcomp>r   )r   r3   r   r   r   r   r   r   r   r   r#   )r   r$   r%   aivpadr   r    r   r!   r0   r"   aes_key_wrap_with_paddingP   s   
r=   wrapped_keyc                    sf  t  dk r
tdt | dvrtdt  dkr@tt| t  }| }| dks1J |d d }|dd  }d}n% fdd	t	d
t  dD }|
d
}	t |}t| |	|\}}d|}tj|dd  dd}
d| |
 }t|d d drd|d  |
  k rd| krn t |d
krt|| d  d| st |d
kr|S |d |  S )Nr'   zMust be at least 16 bytesr&   r*   r   r   r   c                    r,   r-   r!   r.   r>   r!   r"   r1   }   r2   z/aes_key_unwrap_with_padding.<locals>.<listcomp>r   r9   r   r   r8   r:   )r   InvalidUnwrapr3   r   r   r   r5   r   r   r   popr7   r   r   r   r   )r   r>   r%   r5   outr   datar   r   Zencrypted_aivZmlir    r!   r?   r"   aes_key_unwrap_with_paddingi   s6   


" rD   c                    s   t  dk r
tdt  d dkrtdt | dvr tdd} fd	d
tdt  dD }|d}t| ||\}}t||sFt d|S )Nr(   zMust be at least 24 bytesr   r   z-The wrapped key must be a multiple of 8 bytesr&   r*   r+   c                    r,   r-   r!   r.   r?   r!   r"   r1      r2   z"aes_key_unwrap.<locals>.<listcomp>r   )r   r@   r3   r   rA   r7   r   r   )r   r>   r%   r;   r   r   r!   r?   r"   aes_key_unwrap   s   


rE   c                   @   s   e Zd ZdS )r@   N)__name__
__module____qualname__r!   r!   r!   r"   r@      s    r@   )N)typingZ&cryptography.hazmat.primitives.ciphersr   Z1cryptography.hazmat.primitives.ciphers.algorithmsr   Z,cryptography.hazmat.primitives.ciphers.modesr   Z,cryptography.hazmat.primitives.constant_timer   bytesListr#   Anyr4   Tupler7   r=   rD   rE   	Exceptionr@   r!   r!   r!   r"   <module>   s   




1
