
    ~i%                        d dl mZ d dlZd dlZd dlZd dlmZ d dZ	 d!d"dZd#dZ		 d!d$dZ
d%dZdZd&dZd%dZd'dZd(dZ G d d          ZdS ))    )annotationsN)Tuplecert_pemstrreturnTuple[str, str, float]c                T   ddl m} ddlm} |                    |                                  |                      }	 |j                            |j                  j	        
                    |j                  }|rt          |          dk    rnt          j        d|d                   rSt          j        d|d                   r8|d         dd          |d         dd          |j                                        fS n# t"          $ r Y nw xY wd	d	|j                                        fS )
Nr   x509)default_backend   
^ring\..*$	^key\..*$       )cryptographyr   cryptography.hazmat.backendsr   load_pem_x509_certificateencode
extensionsget_extension_for_classSubjectAlternativeNamevalueget_values_for_typeDNSNamelenrematchnot_valid_after_utc	timestamp	Exception)r   r   r   certdnss        v/root/.openclaw/workspace/.venvs/ark-sdk/lib/python3.11/site-packages/volcenginesdkarkruntime/_utils/_key_agreement.pyget_cert_infor&      sB   !!!!!!<<<<<<))(//*;*;__=N=NOODo55'
 

##DL11 	 	PC1A// s1v..  q6!"":s1vabbz4+C+M+M+O+OOO   r4+557777s   B>C> >
D
D    keybytesivplain_bytesassociated_datac                0   ddl m}m}m}  ||                    |           |                    |                                                    }|                    |           |                    |          |	                                z   }||j
        z   S )Nr   Cipher
algorithmsmodes)&cryptography.hazmat.primitives.ciphersr/   r0   r1   AESGCM	encryptorauthenticate_additional_dataupdatefinalizetag)	r(   r*   r+   r,   r/   r0   r1   r5   
ciphertexts	            r%   aes_gcm_encrypt_bytesr;   )   s     QPPPPPPPPPs		"  ikk  **?;;; !!+..1C1C1E1EEJ	%%r'   nonce	plaintextc                    |                                 }t          | ||          }t          j        |                                          S )zXaes_gcm_encrypt_base64_string Encrypt message from base64 string to string using AES-GCM)r   r;   base64	b64encodedecode)r(   r<   r=   r+   cs        r%   aes_gcm_encrypt_base64_stringrC   <   sA    ""$$Kc5+66AA%%'''r'   cipher_bytesc                N   ddl m}m}m} d}|d|          }|| d         }	 ||                    |           |                    ||	                                                    }
|
                    |           |
                    |          |
	                                z   S )zGaes_gcm_decrypt_bytes Decrypt message from bytes to bytes using AES-GCMr   r.      N)
r2   r/   r0   r1   r3   r4   	decryptorr6   r7   r8   )r(   r*   rD   r,   r/   r0   r1   
tag_lengthcipherr9   rG   s              r%   aes_gcm_decrypt_bytesrJ   D   s     QPPPPPPPPPJ,J;,'F

{||
$C s		"c  ikk  **?;;; F##i&8&8&:&:::r'   r:   c                    t          j        |                                          }t          | ||                                          S )N)r?   decodebytesr   rJ   rA   )r(   r<   r:   rD   s       r%   aes_gcm_decrypt_base64_stringrM   [   s<    %j&7&7&9&9::L e\::AACCCr'   zN(?:[A-Za-z0-9+/]{4})*(?:[A-Za-z0-9+/]{2}==|[A-Za-z0-9+/]{3}=|[A-Za-z0-9+/]{4})datac           	     @   t          |          dk     rdS t          dt          |          d          D ]i}	 t          | ||d|dz                      }|dz   t          |          k    r|c S |t          | |||dz   d                   z   c S # t          $ r Y fw xY wdS )z,decrypt_corner_case Decrypt corner case data   r      r   N)r   rangerM   decrypt_corner_caser"   )r(   r<   rN   i	decrypteds        r%   rS   rS   f   s    
4yy2~~r2s4yy!$$  	5UD1q5M+ +I1uD		!!    23tAEFF|LLLLLL 	 	 	D	 s   3B-B
BBc           	         t          j        t          |          }g }|D ][}	 |                    t	          | ||                     (# t
          $ r' |                    t          | ||                     Y Xw xY wd                    |          S )Nr   )r   findallbase64_patternappendrM   r"   rS   join)r(   r<   r:   base64_arrayresultb64s         r%   aes_gcm_decrypt_base64_listr^   u   s    :nj99LF @ @	@MM7UCHHIIII 	@ 	@ 	@MM-c5#>>?????	@776??s   $A.A87A8boolc                    |                                  }t          j        |          }t          |          dz  t          |          dz  cxk    ot          |          dz  dz
  k    nc S )Nr      r   )r   r?   rL   r   )r:   rD   cipher_b64_bytess      r%   decrypt_validaterc      s    $$&&L),77LA%5!6!6	"
 	' 	' 	' 	'"Q&	' 	' 	' 	'r'   c                    t          dg          | j                            dd          z   | j                            dd          z   S )Nr       big)r)   xto_bytesy)r(   s    r%   marshal_cryptography_pub_keyrj      s<     !::r5111CENN2u4M4MMMr'   c                  V    e Z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 )"key_agreement_clientcertificate_pem_stringr   r   Nonec                   d}ddl m} ||k     r#t          d                    ||                    ddl m} ddlm} |                                }|                    |          | _	        | j	        
                                                                }|j        | j	        
                                j        j                 | _        |                    |j        |j        | j                  
                                | _        | j	        j                                        | _        t/          j                    dz   | _        |                                  dS )	z Load cert and extract public keyz42.0.0r   )__version__zThe cryptography package of Ark SDK only supports versions after {}, please install the cryptography package by using pip install "cryptography>={}"r
   eci:	 N)r   rp   r"   formatr   )cryptography.hazmat.primitives.asymmetricrr   r   r   _cert
public_keypublic_numbers_CURVE_TYPEScurvename_curveEllipticCurvePublicNumbersrg   ri   _public_keyr    r!   _not_valid_after_utctime_reload_timeinit_cert_ring_key_id)selfrm   __fixed_version__rp   r   rr   pem_datacert_pubs           r%   __init__zkey_agreement_client.__init__   s[   $,,,,,,***bbhbh%'8c c   	&%%%%%@@@@@@)002233H==
:((**99;;odj&;&;&=&=&C&HI88J
DK
 

*,, 	 %)J$B$L$L$N$N! IKK*;;""$$$$$r'   r=   Tuple[bytes, bytes, str, str]c                `    |                                  \  }}}t          |||          }||||fS )z7encrypt_string encrypt plaintext with ECIES DH protocol)generate_ecies_key_pairrC   )r   r=   r(   r<   tokenr:   s         r%   encrypt_stringz#key_agreement_client.encrypt_string   s:     88::UE23yII
E5*,,r'   r(   r)   r<   c                (    t          |||          }|S )z@encrypt_string_with_key encrypt plaintext with ECIES DH protocol)rC   )r   r(   r<   r=   r:   s        r%   encrypt_string_with_keyz,key_agreement_client.encrypt_string_with_key   s     33yII
r'   r:   c                $    t          |||          S )zAdecrypt_string_with_key decrypt ciphertext with ECIES DH protocol)rM   )r   r(   r<   r:   s       r%   decrypt_string_with_keyz,key_agreement_client.decrypt_string_with_key   s     -S%DDDr'   Tuple[bytes, bytes, str]c                   ddl m} ddlm} ddlm} |                    | j                  }|                    |	                                | j
                  }|                                                                }d} ||                                |dd                              |          }|dd         }	|d|         }
t          |          }|	|
t!          j        |                                          fS )	z/generate_ecies_key_pair generate ECIES key pairr   )hashes)HKDFrq   ,   N)	algorithmlengthsaltinfore   )cryptography.hazmat.primitivesr   'cryptography.hazmat.primitives.kdf.hkdfr   rt   rr   generate_private_keyr{   exchangeECDHr}   rv   rw   SHA256deriverj   r?   r@   rA   )r   r   r   rr   peer_private_keydhRr   bufr(   r<   r   s               r%   r   z,key_agreement_client.generate_ecies_key_pair   s   999999@@@@@@@@@@@@ 224;??&&rwwyy$2BCC''))88:: dmmoo	
 
 

 &** 	 #2#hBvI,Q//E6+E2299;;;;r'   r_   c                :    t          j                     | j        k    S )z)need_reload check if the cert need reload)r   r   r   s    r%   need_reloadz key_agreement_client.need_reload   s    y{{T...r'   c                   	 ddl m} | j        j                            |j                  j                            |j                  }|rut          |          dk    rbt          j        d|d                   rGt          j        d|d                   r,|d         dd         | _        |d         dd         | _        dS d	| _        d	| _        dS # t          $ r Y dS w xY w)
z7init_cert_ring_key_id init ring id and key id from certr   r
   r   r   r   r   Nr   r   )r   r   ru   r   r   r   r   r   r   r   r   r   _ring_id_key_idr"   )r   r   r$   s      r%   r   z*key_agreement_client.init_cert_ring_key_id   s    	))))))*'??+ ''55  
"HHqLLH]CF33 !H\3q622 ! !$Aqrr
"1vabbz "! 	 	 	DD	s   B<C  C 
CCTuple[str, str]c                    | j         | j        fS )z5get_cert_ring_key_id get ring id and key id from cert)r   r   r   s    r%   get_cert_ring_key_idz)key_agreement_client.get_cert_ring_key_id   s    }dl**r'   floatc                    | j         S )z1get_cert_expiration_time get cert expiration time)r~   r   s    r%   get_cert_expiration_timez-key_agreement_client.get_cert_expiration_time   s    ((r'   N)rm   r   r   rn   )r=   r   r   r   r(   r)   r<   r)   r=   r   r   r   r(   r)   r<   r)   r:   r   r   r   )r   r   )r   r_   )r   rn   )r   r   )r   r   )__name__
__module____qualname__r   r   r   r   r   r   r   r   r    r'   r%   rl   rl      s        % % % %4- - - -   E E E E
< < < <2/ / / /   *+ + + +) ) ) ) ) )r'   rl   )r   r   r   r   )r'   )
r(   r)   r*   r)   r+   r)   r,   r)   r   r)   r   )
r(   r)   r*   r)   rD   r)   r,   r)   r   r)   r   )r(   r)   r<   r)   rN   r   r   r   )r:   r   r   r_   )r   r)   )
__future__r   r   r   r?   typingr   r&   r;   rC   rJ   rM   rX   rS   r^   rc   rj   rl   r   r'   r%   <module>r      sa   # " " " " " 				        8 8 8 8, IL& & & & &&( ( ( ( JM; ; ; ; ;.D D D D V 
   	 	 	 	   N N N Ne) e) e) e) e) e) e) e) e) e)r'   