o
    5hE                     @   s   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 d dl	m
Z
 d dlmZ d dlmZ d dlmZmZmZ G d	d
 d
ejZG dd dejZedkrTe  i fddZedkrjdd Zejdd dS dS )    N)	unhexlify)HPKE)DeserializeError)ECC)list_test_cases)DH)SHA256SHA384SHA512c                   @   s   e Zd ZejddZejddZddddddZd	d
 Zdd Z	dd Z
dd Zdd Zdd Zdd Zdd Zdd Zdd Zdd Zdd  Zd!d" Zd#d$ Zd%d& Zd'S )(
HPKE_Testsp256curveA   a          8   )r   p384p521
curve25519curve448c           
      C   s   t j|d}|}tj| |d}| t|j| j|  |j	ddd}tj|||jd}|j
|dd}| d| |	d}|
|}	| d|	 d S )Nr   receiver_keyaead_id   ABC   DEF	auth_datar   r   encs   GHI)r   generater   new
public_keyassertEquallenr    curvessealunseal)
selfr   r   key1	encryptorct	decryptorptct2pt2 r1   v/var/www/html/construction_image-detection-poc/venv/lib/python3.10/site-packages/Crypto/SelfTest/Protocol/test_HPKE.py
round_trip   s    


zHPKE_Tests.round_tripc                 C   s,   | j  D ]}tjD ]}| || q
qd S N)r&   keysr   AEADr3   )r)   r   r   r1   r1   r2   test_round_trip5   s
   
zHPKE_Tests.test_round_tripc                 C   "   t jj}t j| j |dd d S )N)   as    cccccccccccccccccccccccccccccccc)r   r   pskr   r6   
AES128_GCMr"   r*   r#   r)   r   r1   r1   r2   test_psk:   
   
zHPKE_Tests.test_pskc                 C   r8   )Ns   baba)r   r   infor;   r=   r1   r1   r2   	test_info@   r?   zHPKE_Tests.test_infoc                 C   sb   t jdd}| t}tj| tjjd W d    n1 s!w   Y  | 	dt
|j d S )Np224r   r   zUnsupported curve)r   r!   assertRaises
ValueErrorr   r"   r#   r6   r<   assertInstr	exceptionr)   key3cmr1   r1   r2   test_neg_unsupported_curveF   s   
z%HPKE_Tests.test_neg_unsupported_curvec                 C   sX   |  t}tj| j| jtjjd W d    n1 sw   Y  | dt	|j
 d S )Nr   
sender_keyr   zExactly 1 private key)rC   rD   r   r"   r*   key2r6   r<   rE   rF   rG   r)   rJ   r1   r1   r2   test_neg_too_many_private_keysM   s   z)HPKE_Tests.test_neg_too_many_private_keysc                 C   sf   t jdd}| t}tj| j |tjj	d W d    n1 s#w   Y  | 
dt|j d S )Nr   r   rL   zbut recipient key)r   r!   rC   rD   r   r"   r*   r#   r6   r<   rE   rF   rG   rH   r1   r1   r2   test_neg_curve_mismatchT   s   z"HPKE_Tests.test_neg_curve_mismatchc                 C   s   |  t}tj| j dtjjd W d    n1 sw   Y  |  t}tj| j dtjjd W d    n1 s?w   Y  |  t}tj| j dtjjd W d    n1 saw   Y  | dt	|j
 d S )N)    s    GGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGG)r   r:   r   )   JJJrR   )rS   s   YYYYYYYYYYYYYYYYYYYYYYYYYYYYYYYzat least 32)rC   rD   r   r"   r*   r#   r6   r<   rE   rF   rG   rO   r1   r1   r2   test_neg_psk\   s&   zHPKE_Tests.test_neg_pskc                 C   s   d}|  t tj| jtjj|d W d    n1 sw   Y  |  t}tj| j | j j	ddtjjd W d    n1 sFw   Y  | 
dt|j |  t}tj| jtjjd W d    n1 snw   Y  | 
dt|j d S )	NsA   8888888888888888888888888888888888888888888888888888888888888888r   rawformat)r   r    r   z'enc' cannot be an inputr   z'enc' required)rC   r   r   r"   r*   r6   r<   rD   r#   
export_keyrE   rF   rG   )r)   	wrong_encrJ   r1   r1   r2   test_neg_wrong_encm   s(   zHPKE_Tests.test_neg_wrong_encc                 C   s^   t j| jt jj| j jddd}| t	 |
d W d    d S 1 s(w   Y  d S )NrU   rV   r   s<   XYZXYZXYZXYZXYZXYZXYZXYZXYZXYZXYZXYZXYZXYZXYZXYZXYZXYZXYZXYZ)r   r"   r*   r6   CHACHA20_POLY1305rN   r#   rX   rC   rD   r(   )r)   r-   r1   r1   r2   test_neg_unseal_wrong_ct   s   "z#HPKE_Tests.test_neg_unseal_wrong_ctc                 C   sx   t jj}t j| j |d}|jddd}t j| j||jd}| t	 |
| W d    d S 1 s5w   Y  d S )Nr   r   r   r   r   )r   r6   r[   r"   r*   r#   r'   r    rC   rD   r(   )r)   r   r+   r,   r-   r1   r1   r2   test_neg_unseal_no_auth_data   s   "z'HPKE_Tests.test_neg_unseal_no_auth_datac                 C   s   d}t t|}d}t|}d}t|}d}t|}d}	t|	}
d}t|}d}t|}d}t|}tjj}tj||||d	}|||
}| 	|| |||}| 	|| d S )
N@4612c550263fc8ad58375df3f557aac531d26850903e55a9f23f21d8534e8ac8:4265617574792069732074727574682c20747275746820626561757479Zf938558b5d72f1a23810b4be2ab4f84331acc02fc97babc53a52ae8218a355a96d8770ac83d07bea87e13c512a@37fda3567bdbd628e88668c3c8d7e97d1d1253b6d4ea6d44c150f741f1bf4431436f756e742d30436f756e742d31(4f6465206f6e2061204772656369616e2055726eZaf2d7e9ac9ae7e270f46ba1f975be53c09f8d875bdc8535458c2494e8a6eab251c03d0c22a56b8ca42c2063b84)r   r   r@   r    
r   import_x25519_private_keybytesfromhexr   r6   r<   r"   r(   r$   )r)   keyR_hexkeyRpt_hexr.   ct0_hexct0enc_hexr    aad0_hexaad0aad1_hexaad1info_hexr@   ct1_hexct1r   r-   pt_X0pt_X1r1   r1   r2   test_x25519_mode_0   s4   






zHPKE_Tests.test_x25519_mode_0c                 C   s   d}t t|}d}t|}d}t|}d}t|}d}	t|	}
d}t|}d}t|}d}t|}d	}t|}d
}t|}tjj}tj|||||f|d}||
|}| 	|| |||}| 	|| d S )N@c5eb01eb457fe6c6f57577c5413b931550a162c71a03ac8d196babbd4e5ce0fd,456e6e796e20447572696e206172616e204d6f726961@0247fd33b913760fa1fa51e1892d9f307fbe65eb171e8132c2af18555a738b82r_   Ze52c6fed7f758d0cf7145689f21bc1be6ec9ea097fef4e959440012f4feb73fb611b946199e681f4cfc34db8ea@0ad0950d9fb9588e59690b74f1237ecdf1d775cd60be2eca57af5a4b0471c91brb   rc   rd   Z49f3b19b28a9ea9f43e8c71204c00d4a490ee7f61387b6719db765e948123b45b61633ef059ba22cd62437c8ba)r   r   r@   r:   r    rf   )r)   rj   rk   
psk_id_hexpsk_idpsk_hexr:   rl   r.   rm   rn   ro   r    rp   rq   rr   rs   rt   r@   ru   rv   r   r-   rw   rx   r1   r1   r2   test_x25519_mode_1   s>   








zHPKE_Tests.test_x25519_mode_1c                 C   s   d}t t|}d}t t|}d}t|}d}t|}d}	t|	}
d}t|}d}t|}d}t|}d	}t|}tjj}tj|| |||
d
}|	||}| 
|| |	||}| 
|| d S )N@fdea67cf831f1ca98d8e27b1f6abeb5b7745e9d35348b80fa407ff6958f9137e@dc4a146313cce60a278a5323d321f051c5707e9c45ba21a3479fecdf76fc69ddr_   Z5fd92cc9d46dbf8943e72a07e42f363ed5f721212cd90bcfd072bfd9f44e06b80fd17824947496e21b680c141b@23fb952571a14a25e3d678140cd0e5eb47a0961bb18afcf85896e5453c312e76rb   rc   rd   Zd3736bb256c19bfa93d79e8f80b7971262cb7c887e35c26370cfed62254369a1b52e3d505b79dd699f002bc8ed)r   rM   r   r@   r    r   rg   rh   ri   r   r6   r<   r"   r#   r(   r$   )r)   rj   rk   keyS_hexkeySrl   r.   rm   rn   ro   r    rp   rq   rr   rs   rt   r@   ru   rv   r   r-   rw   rx   r1   r1   r2   test_x25519_mode_2   s:   






zHPKE_Tests.test_x25519_mode_2c                 C   s   d}t t|}d}t t|}d}t|}d}t|}d}	t|	}
d}t|}d}t|}d}t|}d	}t|}d
}t|}d}t|}tjj}tj|| |||f||d}|	||}| 
||
 |	||}| 
||
 d S )N@cb29a95649dc5656c2d054c1aa0d3df0493155e9d5da6d7e344ed8b6a64a9423@fc1c87d2f3832adb178b431fce2ac77c7ca2fd680f3406c77b5ecdf818b119f4r{   r|   r_   Za84c64df1e11d8fd11450039d4fe64ff0c8a99fca0bd72c2d4c3e0400bc14a40f27e45e141a24001697737533e@820818d3c23993492cc5623ab437a48a0a7ca3e9639c140fe1e33811eb844b7crb   rc   rd   Z4d19303b848f424fc3c3beca249b2c6de0a34083b8e909b6aa4c3688505c05ffe0c8f57a0a4c5ab9da127435d9)r   rM   r   r:   r@   r    r   )r)   rj   rk   r   r   r   r   r   r:   rl   r.   rm   rn   ro   r    rp   rq   rr   rs   rt   r@   ru   rv   r   r-   rw   rx   r1   r1   r2   test_x25519_mode_3  sD   








zHPKE_Tests.test_x25519_mode_3N)__name__
__module____qualname__r   r!   r*   rN   r&   r3   r7   r>   rA   rK   rP   rQ   rT   rZ   r\   r]   ry   r   r   r   r1   r1   r1   r2   r      s0    (/,r   c                   @   s,   e Zd Zdd Zdd Zdd Zdd Zd	S )
HPKE_TestVectorsc              	   C   s   g | _ z3dd l}tj|j}tj|dd}t|d}t	|| _ W d    W d S 1 s/w   Y  W d S  t
tfyF   td Y d S w )Nr   ProtocolzHPKE-test-vectors.jsonrzN
Warning: skipping extended tests for HPKE (install pycryptodome-test-vectors))vectorspycryptodome_test_vectorsospathdirname__file__joinopenjsonloadFileNotFoundErrorImportErrorprint)r)   r   init_dirfull_file_namefr1   r1   r2   setUpO  s   &zHPKE_TestVectors.setUpc                 C   s   t |}|dkrtjdtj|dddS |dkr$tjdtj|dddS |dkr4tjd	tj|dddS |d
kr=t|S |dkrFt|S d S )N   r   big)	byteorder)r   d   r      r   r   !   )r   r   	constructint
from_bytesr   rg   import_x448_private_key)r)   key_hexkem_idkey_binr1   r1   r2   import_private_keyZ  s$   




z#HPKE_TestVectors.import_private_keyc              	   C   s0  | j s| d t| j D ]\}}|d }|d }|d }|dkr"qtttttd}|||f}|du r6q| j|||dH | |d	 |	 }d}	d
|v rW| |d
 |}	| |d |}
t
j||||	|
\}}| | |d  | |t|d  W d   n1 sw   Y  tdddd qdS )z+Test HPKE encapsulation using test vectors.No test vectors availabler   kdf_idr     )r      )r      )r      )r   r   )r   r   Nidxr   r   skRmskSmskEmr    shared_secret. Tendflush)r   skipTest	enumerater   r	   r
   getsubTestr   r#   r   HPKE_Cipher_encapr$   hexr   r   )r)   r   vectorr   r   r   supported_combihashmodreceiver_pubsender_priv	encap_keyr   r    r1   r1   r2   test_hpke_encapj  sT   



z HPKE_TestVectors.test_hpke_encapc                 C   sn  | j s| d t| j D ]\}}|d }|d }|d }|dkr"qd}||f|vr+q| j|||dr | |d |}d	}d
|v rN| |d
 |}	|	 }t|d }
d	}d|v rft|d t|d f}tj|t	||
||t|d d}|d D ]!}t|d }t|d }t|d }|
||}| ||d q|W d	   n1 sw   Y  tdddd qd	S )z7Test HPKE encryption and decryption using test vectors.r   r   r   r   r   r   r   r   Nr   r    r   r:   r@   )r   r   r    rM   r:   r@   encryptionsr.   r,   aadzDecryption failedr   r   Tr   )r   r   r   r   r   r#   r   r   r"   r6   r(   r$   r   )r)   r   r   r   r   r   r   receiver_priv
sender_pubr   r   r:   receiver_hpke
encryption	plaintext
ciphertextr   	decryptedr1   r1   r2   test_hpke_unseal  sT   



"z!HPKE_TestVectors.test_hpke_unsealN)r   r   r   r   r   r   r   r1   r1   r1   r2   r   M  s
    3r   __main__c                 C   s*   g }|t t7 }| dr|t t7 }|S )N
slow_tests)r   r   r   r   )configtestsr1   r1   r2   	get_tests  s
   
r   c                   C   s   t t S r4   )unittest	TestSuiter   r1   r1   r1   r2   suite  s   r   )defaultTest)r   r   r   binasciir   Crypto.Protocolr   Crypto.Protocol.HPKEr   Crypto.PublicKeyr   Crypto.SelfTest.st_commonr   r   Crypto.Hashr   r	   r
   TestCaser   r   r   mainr   r   r1   r1   r1   r2   <module>   s,      ? 