Ë
    ìmxiÅn  ã                   ó4  — d dl Z d dlmZ d dlZ	 d dlmZ d dlmc m	Z
 d dlmZmZmZmZmZ dZdZdZdZdZe j.                  j0                  e j.                  j                  gZ e j4                  d	¬
«      d„ «       Z e j4                  d	¬
«      d„ «       Z e j4                  d	¬
«      d„ «       Zd„ Zd„ Zd„ Z d„ Z!d„ Z"	 d1d„Z#d„ Z$d„ Z%d„ Z&d„ Z'd„ Z(d„ Z)d„ Z*d„ Z+d„ Z,d„ Z-d„ Z.d „ Z/d!„ Z0e j.                  jc                  d"¬#«      d$„ «       Z2d%„ Z3e j.                  ji                  d& e jj                  ddd'¬(«       e jj                  d)dd*¬(«       e jj                  dd)d+¬(«       e jj                  d)d)d,¬(«      g«      d-„ «       Z6d.„ Z7d/„ Z8d0„ Z9y# e$ r dZdZ
Y �Œiw xY w)2é    N)Ú	timedelta)ÚInMemoryKmsClientÚMockVersioningKmsClientÚverify_file_encryptedÚread_external_keys_to_dictÚparse_wrapped_keyzencrypted_table.in_mem.parquets   0123456789112345Ú
footer_keys   1234567890123450Úcol_keyÚmodule)Úscopec                  óÊ   — t         j                  j                  t        j                  g d¢«      t        j                  g d¢«      t        j                  g d¢«      dœ«      } | S )N)é   é   é   )ÚaÚbÚc)ÚxÚyÚz)ÚpaÚTableÚfrom_pydictÚarray)Ú
data_tables    ú\/home/htdocs/ttos/venv/lib/python3.12/site-packages/pyarrow/tests/parquet/test_encryption.pyr   r   2   sG   € ä—‘×%Ñ%Ü�X‰X’iÓ Ü�X‰X’oÓ&Ü�X‰X’oÓ&ñ'ó €Jð
 Ðó    c                  óL   — t        j                  t        t        ddgi¬«      } | S )Nr   r   )r	   Úcolumn_keys©ÚpeÚEncryptionConfigurationÚFOOTER_KEY_NAMEÚCOL_KEY_NAME)Úbasic_encryption_configs    r   r%   r%   <   s,   € ä ×8Ñ8Ü"ä˜3 ˜*ð
ôÐð
 #Ð"r   c                  óN   — t        j                  t        t        ddgid¬«      } | S )Nr   r   F)r	   r   Úinternal_key_materialr    )Úexternal_encryption_configs    r   r(   r(   F   s1   € ä!#×!;Ñ!;Ü"ä˜3 ˜*ð
ð $ô"%Ðð &Ð%r   c                 óf   — t        j                  | ¬«      }d„ }t        j                  |«      }||fS )z…
    Sets up and returns the KMS connection configuration and crypto factory
    based on provided KMS configuration parameters.
    ©Úcustom_kms_confc                 ó   — t        | «      S ©N©r   ©Úkms_connection_configurations    r   Úkms_factoryz1setup_encryption_environment.<locals>.kms_factoryX   ó   € Ü Ð!=Ó>Ð>r   )r!   ÚKmsConnectionConfigÚCryptoFactory)r+   Úkms_connection_configr1   Úcrypto_factorys       r   Úsetup_encryption_environmentr7   Q   s7   € ô
 ×2Ñ2À?ÔSÐò?ô ×%Ñ% kÓ2€Nà  .Ð0Ð0r   c                 óŒ   — ||j                  d«      ||j                  d«      i}t        |«      \  }}	t        | ||||	«       ||	fS )zL
    Writes an encrypted parquet file based on the provided parameters.
    úUTF-8)Údecoder7   Úwrite_encrypted_parquet)
Úpathr   Úfooter_key_nameÚcol_key_namer	   r
   Úencryption_configr+   r5   r6   s
             r   Úwrite_encrypted_filer@   a   sc   € ð 	˜×*Ñ*¨7Ó3Ø�g—n‘n WÓ-ð€Oô -IØó-Ñ)Ð˜>ô ˜D *Ð.?Ø1°>ôCð ! .Ð0Ð0r   c           	      óv  — | t         z  }t        j                  t        t        ddgidt        d¬«      d¬«      }|j                  du sJ ‚t        ||t        t        t        t        |«      \  }}t        |«       t        j                  t        d¬«      ¬	«      }t        ||||«      }|j                  |«      sJ ‚y
)úDWrite an encrypted parquet, verify it's encrypted, and then read it.r   r   Ú
AES_GCM_V1ç      @©Úminutesé   ©r	   r   Úencryption_algorithmÚcache_lifetimeÚdata_key_length_bitsF©rJ   N)ÚPARQUET_NAMEr!   r"   r#   r$   r   Úuniform_encryptionr@   Ú
FOOTER_KEYÚCOL_KEYr   ÚDecryptionConfigurationÚread_encrypted_parquetÚequals©Útempdirr   r<   r?   r5   r6   Údecryption_configÚresult_tables           r   Ú!test_encrypted_parquet_write_readrX   w   sÅ   € à”\Ñ!€Dô
 ×2Ñ2Ü"ä˜3 ˜*ð
ð *Ü ¨Ô-Ø ô"Ðð ×/Ñ/°5Ñ8Ð8Ð8ä,@Øˆjœ/¬<¼ÄWØó-Ñ)Ð˜>ô ˜$Ôô ×2Ñ2Ü ¨Ô-ô/Ðä)ØÐÐ!6¸óH€Là×Ñ˜\Ô*Ð*Ñ*r   c           	      ó^  — | t         z  }t        j                  t        ddt	        d¬«      d¬«      }|j
                  du sJ ‚t        ||t        t        t        d|«      \  }}t        |«       t        j                  t	        d¬«      ¬«      }t        ||||«      }|j                  |«      sJ ‚y	)
rB   TrC   rD   rE   rG   )r	   rN   rI   rJ   rK   r   rL   N)rM   r!   r"   r#   r   rN   r@   r$   rO   r   rQ   rR   rS   rT   s           r   Ú)test_uniform_encrypted_parquet_write_readrZ   –   s¸   € à”\Ñ!€Dô ×2Ñ2Ü"ØØ)Ü ¨Ô-Ø ô"Ðð ×/Ñ/°4Ñ7Ð7Ð7ä,@Øˆjœ/¬<¼ÀSØó-Ñ)Ð˜>ô ˜$Ôô ×2Ñ2Ü ¨Ô-ô/Ðä)ØÐÐ!6¸óH€Là×Ñ˜\Ô*Ð*Ñ*r   c                 ó   — |j                   r|j                  ||«      }n|j                  ||| «      }|€J ‚t        j                  | |j                  |¬«      5 }|j                  |«       d d d «       y # 1 sw Y   y xY w)N)Úencryption_properties)r'   Úfile_encryption_propertiesÚpqÚParquetWriterÚschemaÚwrite_table)r<   Útabler?   r5   r6   r]   Úwriters          r   r;   r;   ±   s‰   € à×.Ò.Ø%3×%NÑ%NØ!Ð#4ó&6Ñ"ð &4×%NÑ%NØ!Ð#4°dó&<Ð"à%Ð1Ð1Ð1Ü	×	Ñ	Ø�%—,‘,Ø"<ô
>ð "àAGØ×Ñ˜5Ô!÷"÷ "ñ "ús   ÁA4Á4A=Tc                 ó^  — |r|j                  ||«      }n|j                  ||| «      }|€J ‚t        j                  | |¬«      }|j                  dk(  sJ ‚t        j                  | |¬«      }t        |j                  «      dk(  sJ ‚t        j                  | |¬«      }|j                  d¬«      S )N©Údecryption_propertiesr   T©Úuse_threads)	Úfile_decryption_propertiesr^   Úread_metadataÚnum_columnsÚread_schemaÚlenÚnamesÚParquetFileÚread)	r<   rV   r5   r6   r'   ri   Úmetar`   Úresults	            r   rR   rR   À   sÄ   € ñ Ø%3×%NÑ%NØ!Ð#4ó&6Ñ"ð &4×%NÑ%NØ!Ð#4°dó&<Ð"ð &Ð1Ð1Ð1Ü×ÑØÐ$>ô@€Dà×Ñ˜qÒ Ð Ð Ü�^‰^ØÐ$>ô@€Fäˆv�|‰|Ó Ò!Ð!Ð!ä�^‰^ØÐ$>ô@€Fà�;‰; 4ˆ;Ó(Ð(r   c           	      ó
  — | t         z  }t        j                  t        t        ddgidt        d¬«      d¬«      }t        ||t        t        t        t        |«       t        |«       t        t        t        j                  d«      t        t        j                  d«      i«      \  }}t        j                  t        d¬«      ¬	«      }t        j                  t        d
¬«      5  t!        ||||«       ddd«       y# 1 sw Y   yxY w)zYWrite an encrypted parquet, verify it's encrypted,
    and then read it using wrong keys.r   r   rC   rD   rE   rG   rH   r9   rL   zIncorrect master key used©ÚmatchN)rM   r!   r"   r#   r$   r   r@   rO   rP   r   r7   r:   rQ   ÚpytestÚraisesÚ
ValueErrorrR   )rU   r   r<   r?   Úwrong_kms_connection_configÚwrong_crypto_factoryrV   s          r   Ú+test_encrypted_parquet_write_read_wrong_keyr{   ×   sï   € ð ”\Ñ!€Dô
 ×2Ñ2Ü"ä˜3 ˜*ð
ð *Ü ¨Ô-Ø ô"Ðô ˜˜z¬?¼LÜ#¤WÐ.?ôAô ˜$Ôä8TÜœŸ™¨Ó0Ü”j×'Ñ'¨Ó0ðVó 9Ñ5ÐÐ!5ô
 ×2Ñ2Ü ¨Ô-ô/Ðä	�‰”zÐ)EÔ	Fñ "ÜØÐ#Ð%@Ø ô	"÷"÷ "ñ "ús   Ã!C9Ã9Dc                 óÐ   — t        | |«       t        j                  t        d¬«      5  t	        j
                  | t        z  «      j                  «        ddd«       y# 1 sw Y   yxY w)zmWrite an encrypted parquet, verify it's encrypted,
    but then try to read it without decryption properties.úno decryptionrt   N)rX   rv   rw   ÚIOErrorr^   ro   rM   rp   ©rU   r   s     r   Ú0test_encrypted_parquet_read_no_decryption_configr€   ú   sL   € ô & g¨zÔ:ä	�‰”wÐ&6Ô	7ñ 6Ü
�‰�w¤Ñ-Ó.×3Ñ3Ô5÷6÷ 6ñ 6ús   ¨+AÁA%c                 ó´   — t        | |«       t        j                  t        d¬«      5  t	        j
                  | t        z  «       ddd«       y# 1 sw Y   yxY w)zwWrite an encrypted parquet, verify it's encrypted,
    but then try to read its metadata without decryption properties.r}   rt   N)rX   rv   rw   r~   r^   rj   rM   r   s     r   Ú9test_encrypted_parquet_read_metadata_no_decryption_configr‚     sE   € ô & g¨zÔ:ä	�‰”wÐ&6Ô	7ñ 1Ü
×Ñ˜¤<Ñ/Ô0÷1÷ 1ñ 1úó   ¨AÁAc                 ó´   — t        | |«       t        j                  t        d¬«      5  t	        j
                  | t        z  «       ddd«       y# 1 sw Y   yxY w)zuWrite an encrypted parquet, verify it's encrypted,
    but then try to read its schema without decryption properties.r}   rt   N)rX   rv   rw   r~   r^   rl   rM   r   s     r   Ú7test_encrypted_parquet_read_schema_no_decryption_configr…     sC   € ô & g¨zÔ:Ü	�‰”wÐ&6Ô	7ñ /Ü
�‰�w¤Ñ-Ô.÷/÷ /ñ /úrƒ   c           
      óÜ   — | dz  }t        j                  t        ¬«      }t        j                  t
        d¬«      5  t        ||t        t        t        d|«       ddd«       y# 1 sw Y   yxY w)úMWrite an encrypted parquet, but give only footer key,
    without column key.z)encrypted_table_no_col_key.in_mem.parquet©r	   z4Either column_keys or uniform_encryption must be setrt   r   N)	r!   r"   r#   rv   rw   ÚOSErrorr@   r$   rO   ©rU   r   r<   r?   s       r   Ú'test_encrypted_parquet_write_no_col_keyr‹     sn   € ð Ð@Ñ@€Dô ×2Ñ2Ü"ô$Ðô 
�‰”wð%ô
&ñ Aô 	˜T :¬ÄÜ'¨Ð.?ô	A÷	A÷ Añ Aús   »A"Á"A+c           
      óð   — | dz  }t        j                  t        t        ddgid¬«      }t	        j
                  t        d¬«      5  t        ||t        t        t        d|«       d	d	d	«       y	# 1 sw Y   y	xY w)
r‡   z=encrypted_table_col_key_and_uniform_encryption.in_mem.parquetr   r   T)r	   r   rN   z2Cannot set both column_keys and uniform_encryptionrt   r   N)	r!   r"   r#   r$   rv   rw   r‰   r@   rO   rŠ   s       r   Ú;test_encrypted_parquet_write_col_key_and_uniform_encryptionr�   '  s€   € ð ÐTÑT€Dô ×2Ñ2Ü"ä˜3 ˜*ð
ð  ô!Ðô 
�‰”wØRô
Tñ Aô 	˜T :¬ÄÜ'¨Ð.?ô	A÷A÷ Añ Aús   ÁA,Á,A5c                 óè   — | dz  }|}t        j                  «       }d„ }t        j                  |«      }t        j                  t
        d¬«      5  t        |||||«       ddd«       y# 1 sw Y   yxY w)ú<Write an encrypted parquet, but raise KeyError in KmsClient.ú(encrypted_table_kms_error.in_mem.parquetc                 ó   — t        | «      S r-   r.   r/   s    r   r1   z;test_encrypted_parquet_write_kms_error.<locals>.kms_factoryD  s   € ô !Ð!=Ó>Ð>r   r	   rt   N)r!   r3   r4   rv   rw   ÚKeyErrorr;   ©rU   r   r%   r<   r?   r5   r1   r6   s           r   Ú&test_encrypted_parquet_write_kms_errorr”   ;  sy   € ð Ð?Ñ?€DØ/Ðô ×2Ñ2Ó4Ðò?ô
 ×%Ñ% kÓ2€NÜ	�‰”x |Ô	4ñ Gä  jÐ2CØ 5°~ô	G÷G÷ Gñ Gúó   ÁA(Á(A1c                 ó   ‡— | dz  }|}t        j                  «       } G d„ dt         j                  «      Šˆfd„}t        j                  |«      }t	        j
                  t        d¬«      5  t        |||||«       ddd«       y# 1 sw Y   yxY w)r�   r�   c                   ó"   — e Zd ZdZd„ Zd„ Zd„ Zy)úJtest_encrypted_parquet_write_kms_specific_error.<locals>.ThrowingKmsClientzVA KmsClient implementation that throws exception in
        wrap/unwrap calls
        c                 óP   — t         j                  j                  | «       || _        y)z%Create an InMemoryKmsClient instance.N)r!   Ú	KmsClientÚ__init__Úconfig©Úselfrœ   s     r   r›   zStest_encrypted_parquet_write_kms_specific_error.<locals>.ThrowingKmsClient.__init__^  s   € ä�L‰L×!Ñ! $Ô'Ø ˆD�Kr   c                 ó   — t        d«      ‚)NúCannot Wrap Key©rx   ©rž   Ú	key_bytesÚmaster_key_identifiers      r   Úwrap_keyzStest_encrypted_parquet_write_kms_specific_error.<locals>.ThrowingKmsClient.wrap_keyc  s   € ÜÐ.Ó/Ð/r   c                 ó   — t        d«      ‚)NzCannot Unwrap Keyr¡   ©rž   Úwrapped_keyr¤   s      r   Ú
unwrap_keyzUtest_encrypted_parquet_write_kms_specific_error.<locals>.ThrowingKmsClient.unwrap_keyf  s   € ÜÐ0Ó1Ð1r   N©Ú__name__Ú
__module__Ú__qualname__Ú__doc__r›   r¥   r©   © r   r   ÚThrowingKmsClientr˜   Y  s   „ ñ	ò	!ò
	0ó	2r   r°   c                 ó   •—  ‰| «      S r-   r¯   )r0   r°   s    €r   r1   zDtest_encrypted_parquet_write_kms_specific_error.<locals>.kms_factoryi  s   ø€ á Ð!=Ó>Ð>r   r    rt   N)r!   r3   rš   r4   rv   rw   rx   r;   )	rU   r   r%   r<   r?   r5   r1   r6   r°   s	           @r   Ú/test_encrypted_parquet_write_kms_specific_errorr²   P  s‹   ø€ ð Ð?Ñ?€DØ/Ðô ×2Ñ2Ó4Ðô2œBŸL™Lô 2ô ?ô ×%Ñ% kÓ2€NÜ	�‰”zÐ):Ô	;ñ Gä  jÐ2CØ 5°~ô	G÷G÷ Gñ Gús   Á+BÂBc                 óè   — | dz  }|}t        j                  «       }d„ }t        j                  |«      }t        j                  t
        d¬«      5  t        |||||«       ddd«       y# 1 sw Y   yxY w)z@Write an encrypted parquet, but raise ValueError in kms_factory.ú0encrypted_table_kms_factory_error.in_mem.parquetc                 ó   — t        d«      ‚)NúCannot create KmsClientr¡   r/   s    r   r1   zCtest_encrypted_parquet_write_kms_factory_error.<locals>.kms_factory}  s   € ÜÐ2Ó3Ð3r   r¶   rt   N)r!   r3   r4   rv   rw   rx   r;   r“   s           r   Ú.test_encrypted_parquet_write_kms_factory_errorr·   t  s~   € ð ÐGÑG€DØ/Ðô ×2Ñ2Ó4Ðò4ô ×%Ñ% kÓ2€NÜ	�‰”zØ6ô
8ñ Gô 	   jÐ2CØ 5°~ô	G÷G÷ Gñ Gúr•   c                 óþ   ‡— | dz  }|}t        j                  «       } G d„ d«      Šˆfd„}t        j                  |«      }t        j                  t
        «      5  t        |||||«       ddd«       y# 1 sw Y   yxY w)z_Write an encrypted parquet, but use wrong KMS client type
    that doesn't implement KmsClient.r´   c                   ó"   — e Zd ZdZd„ Zd„ Zd„ Zy)úOtest_encrypted_parquet_write_kms_factory_type_error.<locals>.WrongTypeKmsClientz4This is not an implementation of KmsClient.
        c                 ó&   — |j                   | _        y r-   )r+   Úmaster_keys_mapr�   s     r   r›   zXtest_encrypted_parquet_write_kms_factory_type_error.<locals>.WrongTypeKmsClient.__init__–  s   € Ø#)×#9Ñ#9ˆDÕ r   c                  ó   — y r-   r¯   r¢   s      r   r¥   zXtest_encrypted_parquet_write_kms_factory_type_error.<locals>.WrongTypeKmsClient.wrap_key™  ó   € Ør   c                  ó   — y r-   r¯   r§   s      r   r©   zZtest_encrypted_parquet_write_kms_factory_type_error.<locals>.WrongTypeKmsClient.unwrap_keyœ  r¾   r   Nrª   r¯   r   r   ÚWrongTypeKmsClientrº   ’  s   „ ñ	ò	:ò	ó	r   rÀ   c                 ó   •—  ‰| «      S r-   r¯   )r0   rÀ   s    €r   r1   zHtest_encrypted_parquet_write_kms_factory_type_error.<locals>.kms_factoryŸ  s   ø€ Ù!Ð">Ó?Ð?r   N)r!   r3   r4   rv   rw   Ú	TypeErrorr;   )	rU   r   r%   r<   r?   r5   r1   r6   rÀ   s	           @r   Ú3test_encrypted_parquet_write_kms_factory_type_errorrÃ   ˆ  sƒ   ø€ ð ÐGÑG€DØ/Ðô ×2Ñ2Ó4Ð÷ñ ô@ô ×%Ñ% kÓ2€NÜ	�‰”yÓ	!ñ Gä  jÐ2CØ 5°~ô	G÷G÷ Gñ Gús   ÁA3Á3A<c            
      óJ  — d„ } t        j                  t        t        ddgidddt	        d¬«      dd	¬
«      } | |«       t        j                  t        ¬«      }t        ddgi|_        d|_        d|_        d|_        t	        d¬«      |_	        d|_
        d	|_         | |«       y )Nc                 ó0  — t         | j                  k(  sJ ‚ddg| j                  t           k(  sJ ‚d| j                  k(  sJ ‚| j
                  sJ ‚| j                  rJ ‚t        d¬«      | j                  k(  sJ ‚| j                  rJ ‚d| j                  k(  sJ ‚y )Nr   r   ÚAES_GCM_CTR_V1ç      $@rE   éÀ   )r#   r	   r   r$   rI   Úplaintext_footerÚdouble_wrappingr   rJ   r'   rK   )r?   s    r   Ú!validate_encryption_configurationzZtest_encrypted_parquet_encryption_configuration.<locals>.validate_encryption_configurationª  s¦   € ÜÐ"3×">Ñ">Ò>Ð>Ð>Ø�SˆzÐ.×:Ñ:¼<ÑHÒHÐHÐHØÐ#4×#IÑ#IÒIÐIÐIØ ×1Ò1Ð1Ð1Ø$×4Ò4Ð4Ð4Ü Ô&Ð*;×*JÑ*JÒJÐJÐJØ$×:Ò:Ð:Ð:ØÐ'×<Ñ<Ò<Ð<Ñ<r   r   r   rÆ   TFrÇ   rE   rÈ   )r	   r   rI   rÉ   rÊ   rJ   r'   rK   rˆ   )r!   r"   r#   r$   r   r   rI   rÉ   rÊ   rJ   r'   rK   )rË   r?   Úencryption_config_1s      r   Ú/test_encrypted_parquet_encryption_configurationrÍ   ©  sµ   € ò=ô ×2Ñ2Ü"Ü! C¨ :Ð0Ø-ØØÜ ¨Ô.Ø#Ø ô	Ðñ &Ð&7Ô8ä×4Ñ4Ü"ô$Ðä'3°c¸3°ZÐ&BÐÔ#Ø/?ÐÔ,Ø+/ÐÔ(Ø*/ÐÔ'Ü)2¸4Ô)@ÐÔ&Ø05ÐÔ-Ø/2ÐÔ,Ù%Ð&9Õ:r   c                  óú   — t        j                  t        d¬«      ¬«      } t        d¬«      | j                  k(  sJ ‚t        j                  «       }t        d¬«      |_        t        d¬«      |j                  k(  sJ ‚y )NrÇ   rE   rL   )r!   rQ   r   rJ   )rV   Údecryption_config_1s     r   Ú/test_encrypted_parquet_decryption_configurationrÐ   Ì  sl   € Ü×2Ñ2Ü ¨Ô.ô0Ðä˜TÔ"Ð&7×&FÑ&FÒFÐFÐFä×4Ñ4Ó6ÐÜ)2¸4Ô)@ÐÔ&Ü˜TÔ"Ð&9×&HÑ&HÒHÐHÑHr   c                  óÈ   — d„ } t        j                  ddddddœ¬«      } | |«       t        j                  «       }d|_        d|_        d|_        dddœ|_         | |«       y )	Nc                 ó’   — d| j                   k(  sJ ‚d| j                  k(  sJ ‚d| j                  k(  sJ ‚dddœ| j                  k(  sJ ‚y )NÚ	Instance1ÚURL1ÚMyTokenÚkey_material_1Úkey_material_2©Úkey1Úkey2©Úkms_instance_idÚkms_instance_urlÚkey_access_tokenr+   )r5   s    r   Úvalidate_kms_connection_configzPtest_encrypted_parquet_kms_configuration.<locals>.validate_kms_connection_config×  sb   € ØÐ3×CÑCÒCÐCÐCØÐ.×?Ñ?Ò?Ð?Ð?ØÐ1×BÑBÒBÐBÐBØ)Ð3CÑDØ%×5Ñ5ò6ð 	7ñ 6r   rÓ   rÔ   rÕ   rÖ   r×   rØ   rÛ   )r!   r3   rÜ   rÝ   rÞ   r+   )rß   r5   Úkms_connection_config_1s      r   Ú(test_encrypted_parquet_kms_configurationrá   Ö  s   € ò7ô ×2Ñ2Ø#ØØ"à$Ø$ñ
ô	Ðñ #Ð#8Ô9ä ×4Ñ4Ó6ÐØ.9ÐÔ+Ø/5ÐÔ,Ø/8ÐÔ,à Ø ñ/ÐÔ+ñ #Ð#:Õ;r   zNPlaintext footer - reading plaintext column subset reads encrypted columns too)Úreasonc                 ó>  — | t         z  }t        j                  t        t        ddgidd¬«      }t        j
                  t        t        j                  d«      t        t        j                  d«      i¬«      }d„ }t        j                  |«      }t        |||||«       y	)
z‹Write an encrypted parquet, with plaintext footer
    and with single wrapping,
    verify it's encrypted, and then read plaintext columns.r   r   TF)r	   r   rÉ   rÊ   r9   r*   c                 ó   — t        | «      S r-   r.   r/   s    r   r1   zStest_encrypted_parquet_write_read_plain_footer_single_wrapping.<locals>.kms_factory  r2   r   N)rM   r!   r"   r#   r$   r3   rO   r:   rP   r4   r;   )rU   r   r<   r?   r5   r1   r6   s          r   Ú>test_encrypted_parquet_write_read_plain_footer_single_wrappingrå   ó  sœ   € ð ”\Ñ!€Dô
 ×2Ñ2Ü"ä˜3 ˜*ð
ð ØôÐô ×2Ñ2äœZ×.Ñ.¨wÓ7Üœ'Ÿ.™.¨Ó1ð
ôÐò?ô ×%Ñ% kÓ2€Nä˜D *Ð.?Ø1°>õCr   c           	      ó  — | t         z  }t        ||t        t        t        t
        |«      \  }}t        |«       t        j                  «       }t        ||||d¬«      }t        j                  j                  j                  |«      }t        |j                  «       x}	«      t        |j                   t           «      dz   k(  sJ ‚t#        |	D �
cg c]  }
|j%                  |
«      du‘Œ c}
«      sJ ‚|j'                  |«      sJ ‚yc c}
w )zˆWrite an encrypted parquet file with external key material, verify
    it's encrypted, then read both the table and external store.
    F©r'   r   N)rM   r@   r#   r$   rO   rP   r   r!   rQ   rR   r   Ú_parquet_encryptionÚFileSystemKeyMaterialStoreÚfor_filerm   Úget_key_id_setr   ÚallÚget_key_materialrS   )rU   r   r(   r<   r5   r6   rV   rW   ÚstoreÚkey_idsÚks              r   Ú*test_encrypted_parquet_write_read_externalrñ     s  € ð
 ”\Ñ!€Dä,@Øˆjœ/¬<¼ÄWØ"ó-$Ñ)Ð˜>ô ˜$Ôä×2Ñ2Ó4ÐÜ)ØÐÐ!6¸Ø#ô%€Lô ×"Ñ"×=Ñ=×FÑFÀtÓL€Eä˜%×.Ñ.Ó0Ð0ˆwÓ1ÜÐ&×2Ñ2´<Ñ@ÓAÀAÑEòGð Gð Gä¸wÖG¸!�×&Ñ& qÓ)°Ò5ÒGÔHÐHÐHØ×Ñ˜\Ô*Ð*Ñ*ùò Hs   ÃD)Údouble_wrap_initialÚdouble_wrap_rotatedzdouble wrapping)ÚidFzsingle to double wrappedzdouble to singe wrappedzsingle wrappingc                 óz  ‡‡‡‡— | t         z  }t        j                  t        t        ddgid‰¬«      }t        j
                  d¬«      }d„ }t        j                  |«      }t        |||||«       t        |«      Š|j                  d«       |j                  ||‰¬	«       t        |«      Št        |«       t        |t        j                  «       ||d¬
«      }	t        ‰v sJ ‚t        ‰v sJ ‚t        ‰v sJ ‚t        ‰v sJ ‚dt        ddfˆˆˆˆfd„}
 |
t        «        |
t        «       |j                  |	«      sJ ‚y)a£  Tests CryptoFactory.rotate_master_keys

    Note: The CryptoFactory.rotate_master_keys() double_wrapping keword arg
    may be either True (the default) or False regardless of whether
    EncryptionConfig.double_wrapping was set to true (also the default) when
    the external key material store was written. This means double wrapping may
    be set one way initially and then applied or removed during rotation.
    r   r   F)r	   r   r'   rÊ   Ú1)rÞ   c                 ó   — t        | «      S r-   )r   r/   s    r   r1   z8test_external_key_material_rotation.<locals>.kms_factoryR  s   € Ü&Ð'CÓDÐDr   Ú2)rÊ   rç   Úmaster_key_idÚreturnNc                 óÐ   •— ‰	|    }‰
r|j                   }n|j                  }t        |«      \  }}}‰|    }‰r|j                   }n|j                  }t        |«      \  }}}||k  sJ ‚y r-   )Úwrapped_kekÚwrapped_dekr   )rù   Úbefore_key_matÚbefore_key_wrappedÚ_Ú
before_verÚafter_key_matÚafter_key_wrappedÚ	after_verÚ
after_keysÚbefore_keysrò   ró   s           €€€€r   Úcheck_rotated_external_keyszHtest_external_key_material_rotation.<locals>.check_rotated_external_keysr  s€   ø€ Ø$ ]Ñ3ˆÙØ!/×!;Ñ!;Ñà!/×!;Ñ!;ÐÜ,Ð-?Ó@Ñˆˆ:�qà" =Ñ1ˆÙØ -× 9Ñ 9Ñà -× 9Ñ 9ÐÜ+Ð,=Ó>‰ˆˆ9�að ˜IÒ%Ð%Ñ%r   )rM   r!   r"   r#   r$   r3   r4   r;   r   Úrefresh_key_access_tokenÚrotate_master_keysr   rR   rQ   ÚstrrS   )Úreusable_tempdirr   rò   ró   r<   r?   r5   r1   r6   Útable_read_after_rotationr  r  r  s     ``       @@r   Ú#test_external_key_material_rotationr  5  s]  û€ ð& œlÑ*€DÜ×2Ñ2Ü"Ü! C¨ :Ð.Ø#Ø+ô	-Ðô ×2Ñ2ÀCÔHÐòEä×%Ñ% kÓ2€NÜØØØØØôô -¨TÓ2€Kð ×2Ñ2°3Ô7à×%Ñ%ØØØ+ð &ô -ô
 ,¨DÓ1€JÜ˜$ÔÜ 6ØÜ
×"Ñ"Ó$ØØØ#ô!%Ðô ˜kÑ)Ð)Ð)Ü˜;Ñ&Ð&Ð&Ü˜jÑ(Ð(Ð(Ü˜:Ñ%Ð%Ð%ð&´3ð &¸4÷ &ð &ñ"  ¤Ô0Ù¤Ô-Ø×ÑÐ6Ô7Ð7Ñ7r   c           	      óx  — | t         z  }t        ||t        t        t        t
        |«      \  }}t        |«       t        j                  t        d¬«      ¬«      }t        d«      D ]T  }|j                  ||«      }|€J ‚t        j                  ||¬«      }	|	j                  d¬«      }
|j                  |
«      rŒTJ ‚ y)	z`Write an encrypted parquet, verify it's encrypted,
    and then read it multithreaded in a loop.rD   rE   rL   é2   Nre   Trg   )rM   r@   r#   r$   rO   rP   r   r!   rQ   r   Úrangeri   r^   ro   rp   rS   )rU   r   r%   r<   r5   r6   rV   Úiri   rr   rW   s              r   Útest_encrypted_parquet_loopr  ˆ  sÅ   € ð ”\Ñ!€Dô
 -AØˆjœ/¬<¼ÄWØó-!Ñ)Ð˜>ô ˜$Ôä×2Ñ2Ü ¨Ô-ô/Ðô �2‹Yò 	/ˆà%3×%NÑ%NØ!Ð#4ó&6Ð"à)Ð5Ð5Ð5ä—‘ØÐ(BôDˆà—{‘{¨t�{Ó4ˆØ× Ñ  Õ.Ð.Ð.ñ	/r   c           	      óR  — | t         z  }t        ||t        t        t        t
        |«      \  }}t        |«       t        j                  t        d¬«      ¬«      }|j                  ||«      }~t        j                  ||¬«      }|j                  d¬«      }	|j                  |	«      sJ ‚y)z^
    Test that decryption properties can be used if the crypto factory is no longer alive
    rD   rE   rL   re   Trg   N)rM   r@   r#   r$   rO   rP   r   r!   rQ   r   ri   r^   ro   rp   rS   )
rU   r   r%   r<   r5   r6   rV   ri   rr   rW   s
             r   Ú%test_read_with_deleted_crypto_factoryr  ¥  s¦   € ð ”\Ñ!€DÜ,@Øˆjœ/¬<¼ÄWØó-!Ñ)Ð˜>ô ˜$Ôô ×2Ñ2Ü ¨Ô-ô/Ðà!/×!JÑ!JØÐ0ó"2Ðàä�^‰^ØÐ$>ô@€Fà—;‘;¨4�;Ó0€LØ×Ñ˜\Ô*Ð*Ñ*r   c           	      ój  — | t         z  }t        ||t        t        t        t
        |«      \  }}t        j                  t        d¬«      ¬«      }|j                  ||«      }t        j                  ||¬«      }|j                  |«      sJ ‚t        j                  | |¬«      }|j                  |«      sJ ‚y)z>Write an encrypted parquet then read it back using read_table.rD   rE   rL   re   N)rM   r@   r#   r$   rO   rP   r!   rQ   r   ri   r^   Ú
read_tablerS   )	rU   r   r%   r<   r5   r6   rV   ri   rW   s	            r   Ú!test_encrypted_parquet_read_tabler  ½  sµ   € à”\Ñ!€Dô -AØˆjœ/¬<¼ÄWØó-!Ñ)Ð˜>ô ×2Ñ2Ü ¨Ô-ô/Ðà!/×!JÑ!JØÐ0ó"2Ðô —=‘= Ð=WÔX€Lð ×Ñ˜\Ô*Ð*Ð*ô —=‘=ØÐ'AôC€Là×Ñ˜\Ô*Ð*Ñ*r   )T):rv   Údatetimer   Úpyarrowr   Úpyarrow.parquetÚparquetr^   Úpyarrow.parquet.encryptionÚ
encryptionr!   Ú pyarrow.tests.parquet.encryptionr   r   r   r   r   ÚImportErrorrM   rO   r#   rP   r$   ÚmarkÚparquet_encryptionÚ
pytestmarkÚfixturer   r%   r(   r7   r@   rX   rZ   r;   rR   r{   r€   r‚   r…   r‹   r�   r”   r²   r·   rÃ   rÍ   rÐ   rá   Úxfailrå   rñ   ÚparametrizeÚparamr  r  r  r  r¯   r   r   ú<module>r'     s  ðó" Ý Û ðEÝ ß+Ð+÷
Eõ Eð 0€Ø €
Ø€Ø
€Ø€ð ‡K�K×"Ñ"Ø
‡K�K×Ñð€
ð €‡��hÔñó  ðð €‡��hÔñ#ó  ð#ð €‡��hÔñ&ó  ð&ò1ò 1ò,+ò>+ò6"ð" 26ó)ò. "òF6ò1ò/òAò"Aò(Gò*!GòHGò(GòB ;òFIò<ð: ‡�×Ñð 2Ðó 3ñCó3ðCòN+ð2 ‡�×ÑØ2Øˆ�‰�T˜4Ð$5Ô6Øˆ�‰�U˜DÐ%?Ô@Øˆ�‰�T˜5Ð%>Ô?Øˆ�‰�U˜EÐ&7Ô8ð	5:ó;ñJ8ó;ðJ8òZ/ò:+ó0+øðM ò Ø	€BØ	ƒBðús   �F
 Æ
	FÆF