o
    _bV                     @   s  d dl Z d dlZd dlmZ d dlZd dlZd dlZd dl	m
Z
mZ d dlmZmZmZmZmZ d dlmZmZ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 d dl m!Z! d dl"m#Z#m$Z$ d dl%m&Z&m'Z'm(Z( e )e*Z+G dd dZ,ddde d de ddfddZ-G dd dZ.G dd de$Z/G dd de#Z0G dd dZ1G dd de1Z2G dd  d eZ3G d!d" d"Z4G d#d$ d$Z5G d%d& d&Z6G d'd( d(Z7dS ))    N)BytesIO)AwsCredentialsAwsCredentialsProvider)ClientBootstrapClientTlsContextDefaultHostResolverEventLoopGroupTlsContextOptions)S3ClientS3RequestTlsModeS3RequestType)UNSIGNED)urlsplit)Config)NoCredentialsError)GBMB)TransferNotDoneError)BaseTransferFutureBaseTransferMeta)CallArgsOSUtilsget_callbacksc                   @   s$   e Zd Zdd Zdd Zdd ZdS )CRTCredentialProviderAdapterc                 C   s   || _ d | _t | _d S N)_botocore_credential_provider_loaded_credentials	threadingLock_lock)selfbotocore_credential_provider r"   2usr/lib/python3.10/site-packages/s3transfer/crt.py__init__+   s   z%CRTCredentialProviderAdapter.__init__c                 C   s   |    }t|j|j|jS r   )_get_credentialsZget_frozen_credentialsr   Z
access_keyZ
secret_keytoken)r    Zcredentialsr"   r"   r#   __call__0   s   z%CRTCredentialProviderAdapter.__call__c                 C   sX   | j  | jd u r| j }|d u rt || _| jW  d    S 1 s%w   Y  d S r   )r   r   r   Zload_credentialsr   )r    Zloaded_credsr"   r"   r#   r%   6   s   
$z-CRTCredentialProviderAdapter._get_credentialsN)__name__
__module____qualname__r$   r'   r%   r"   r"   r"   r#   r   *   s    r         Tc              	   C   s   t |}t|}t||}	d}
d}|rtjntj}|dur4t }|r)|j|d nd|_t	|}|
 }|r?t|}t|}
|d t }t|	| |
||||dS )a  
    :type region: str
    :param region: The region used for signing

    :type botocore_credential_provider:
        Optional[botocore.credentials.CredentialResolver]
    :param botocore_credential_provider: Provide credentials for CRT
        to sign the request if not set, the request will not be signed

    :type num_threads: Optional[int]
    :param num_threads: Number of worker threads generated. Default
        is the number of processors in the machine.

    :type target_throughput: Optional[int]
    :param target_throughput: Throughput target in Bytes.
        Default is 0.625 GB/s (which translates to 5 Gb/s).

    :type part_size: Optional[int]
    :param part_size: Size, in Bytes, of parts that files will be downloaded
        or uploaded in.

    :type use_ssl: boolean
    :param use_ssl: Whether or not to use SSL.  By default, SSL is used.
        Note that not all services support non-ssl connections.

    :type verify: Optional[boolean/string]
    :param verify: Whether or not to verify SSL certificates.
        By default SSL certificates are verified.  You can provide the
        following values:

        * False - do not validate SSL certificates.  SSL will still be
            used (unless use_ssl is False), but SSL certificates
            will not be verified.
        * path/to/cert/bundle.pem - A filename of the CA cert bundle to
            use. Specify this argument if you want to use a custom CA cert
            bundle instead of the default one on your system.
    N)Zca_filepathFr,   )	bootstrapregionZcredential_provider	part_sizetls_modetls_connection_optionsZthroughput_target_gbps)r   r   r   r   ZENABLEDZDISABLEDr	   Z&override_default_trust_store_from_pathZverify_peerr   Znew_connection_optionsr   r   Znew_delegater   r
   )r.   r!   Znum_threadsZtarget_throughputr/   Zuse_sslZverifyZevent_loop_groupZhost_resolverr-   Zproviderr1   r0   Ztls_ctx_optionsZclient_tls_optionZcredentails_provider_adapterZtarget_gbpsr"   r"   r#   create_s3_crt_clientB   sB   /
r2   c                   @   s   e Zd ZdddZdd Zdd Z	ddd	Zdd
dZdddZdddZ	dd Z
dd Zdd ZdddZdd Zdd ZdS ) CRTTransferManagerNc                 C   s@   |du rt  | _|| _t|| j| _g | _td| _d| _	dS )aF  A transfer manager interface for Amazon S3 on CRT s3 client.

        :type crt_s3_client: awscrt.s3.S3Client
        :param crt_s3_client: The CRT s3 client, handling all the
            HTTP requests and functions under then hood

        :type crt_request_serializer: s3transfer.crt.BaseCRTRequestSerializer
        :param crt_request_serializer: Serializer, generates unsigned crt HTTP
            request.

        :type osutil: s3transfer.utils.OSUtils
        :param osutil: OSUtils object to use for os-related behavior when
            using with transfer manager.
        N   r   )
r   _osutil_crt_s3_clientS3ClientArgsCreator_s3_args_creator_future_coordinatorsr   	Semaphore
_semaphore_id_counter)r    Zcrt_s3_clientcrt_request_serializerosutilr"   r"   r#   r$      s   
zCRTTransferManager.__init__c                 C   s   | S r   r"   r    r"   r"   r#   	__enter__   s   zCRTTransferManager.__enter__c                 G   s   d}|rd}|  | d S )NFT	_shutdown)r    exc_type	exc_valueargscancelr"   r"   r#   __exit__   s   zCRTTransferManager.__exit__c                 C   s6   |d u ri }|d u ri }t |||||d}| d|S )Nbucketkeyfileobj
extra_argssubscribersZ
get_objectr   _submit_transfer)r    rI   rJ   rK   rL   rM   callargsr"   r"   r#   download   s   zCRTTransferManager.downloadc                 C   s6   |d u ri }|d u ri }t |||||d}| d|S )NrH   Z
put_objectrN   )r    rK   rI   rJ   rL   rM   rP   r"   r"   r#   upload   s   zCRTTransferManager.uploadc                 C   s4   |d u ri }|d u ri }t ||||d}| d|S )N)rI   rJ   rL   rM   Zdelete_objectrN   )r    rI   rJ   rL   rM   rP   r"   r"   r#   delete   s   zCRTTransferManager.deleteFc                 C   s   |  | d S r   rA   r    rF   r"   r"   r#   shutdown      zCRTTransferManager.shutdownc                 C   s    | j D ]
}| s|  qd S r   )r9   donerF   r    coordinatorr"   r"   r#   _cancel_transfers   s
   
z$CRTTransferManager._cancel_transfersc                 C      | j D ]}|  qd S r   )r9   resultrX   r"   r"   r#   _finish_transfers      

z$CRTTransferManager._finish_transfersc                 C   r[   r   )r9   %wait_until_on_done_callbacks_completerX   r"   r"   r#   _wait_transfers_done   r^   z'CRTTransferManager._wait_transfers_donec                 C   sz   |r|    z1z|   W n ty   |    Y n ty"   Y nw W |   d S W |   d S W |   d S |   w r   )rZ   r]   KeyboardInterrupt	Exceptionr`   rT   r"   r"   r#   rB      s   zCRTTransferManager._shutdownc                 K      | j   d S r   )r;   releaser    kwargsr"   r"   r#   _release_semaphore  rV   z%CRTTransferManager._release_semaphorec              
   C   s  | j g}t| jd}t| j||d}td	i |}t|}|| z$| j  | j	
|d}|  | j	|||||}	| jjd	i |	}
W n& tym } z||d | j	j
|d|d}||d W Y d }~n
d }~ww ||
 | j| |  jd7  _|S )
N)transfer_id)metarY   ZqueuedTrW   )after_subscribers)error   r"   )rg   CRTTransferCoordinatorr<   CRTTransferMetaCRTTransferFutureAfterDoneHandlerappendr;   acquirer8   get_crt_callbackget_make_request_argsr6   Zmake_requestrb   set_exceptionset_s3_requestr9   )r    request_type	call_argson_done_after_callsrY   
componentsfutureZ	afterdoneZ	on_queuedZcrt_callargsZcrt_s3_requesteon_doner"   r"   r#   rO     sD   



z#CRTTransferManager._submit_transferr   NNF)r(   r)   r*   r$   r@   rG   rQ   rR   rS   rU   rZ   r]   r`   rB   rg   rO   r"   r"   r"   r#   r3      s    





r3   c                   @   s>   e Zd ZdZdddZedd Zedd Zed	d
 ZdS )rn   z*Holds metadata about the CRTTransferFutureNc                 C   s   || _ || _i | _d S r   )_transfer_id
_call_args_user_context)r    rh   rx   r"   r"   r#   r$   1  s   
zCRTTransferMeta.__init__c                 C      | j S r   )r   r?   r"   r"   r#   rx   6     zCRTTransferMeta.call_argsc                 C   r   r   )r   r?   r"   r"   r#   rh   :  r   zCRTTransferMeta.transfer_idc                 C   r   r   )r   r?   r"   r"   r#   user_context>  r   zCRTTransferMeta.user_contextr~   )	r(   r)   r*   __doc__r$   propertyrx   rh   r   r"   r"   r"   r#   rn   .  s    


rn   c                   @   sD   e Zd ZdddZedd Zdd Zddd	Zd
d Zdd Z	dS )ro   Nc                 C   s    || _ |du rt | _ || _dS )aZ  The future associated to a submitted transfer request via CRT S3 client

        :type meta: s3transfer.crt.CRTTransferMeta
        :param meta: The metadata associated to the transfer future.

        :type coordinator: s3transfer.crt.CRTTransferCoordinator
        :param coordinator: The coordinator associated to the transfer future.
        N)_metarn   _coordinator)r    ri   rY   r"   r"   r#   r$   D  s   	
zCRTTransferFuture.__init__c                 C   r   r   )r   r?   r"   r"   r#   ri   R  r   zCRTTransferFuture.metac                 C   s
   | j  S r   )r   rW   r?   r"   r"   r#   rW   V     
zCRTTransferFuture.donec                 C      | j | d S r   )r   r\   r    timeoutr"   r"   r#   r\   Y     zCRTTransferFuture.resultc                 C   rc   r   )r   rF   r?   r"   r"   r#   rF   \  rV   zCRTTransferFuture.cancelc                 C   s$   |   std| jj|dd dS )z!Sets the exception on the future.z?set_exception can only be called once the transfer is complete.T)overrideN)rW   r   r   ru   )r    	exceptionr"   r"   r#   ru   _  s
   zCRTTransferFuture.set_exceptionr~   r   )
r(   r)   r*   r$   r   ri   rW   r\   rF   ru   r"   r"   r"   r#   ro   C  s    


ro   c                   @   s   e Zd Zdd ZdS )BaseCRTRequestSerializerc                 C   s   t d)al  Serialize CRT HTTP requests.

        :type transfer_type: string
        :param transfer_type: the type of transfer made,
            e.g 'put_object', 'get_object', 'delete_object'

        :type future: s3transfer.crt.CRTTransferFuture

        :rtype: awscrt.http.HttpRequest
        :returns: An unsigned HTTP request to be used for the CRT S3 client
        zserialize_http_request())NotImplementedError)r    transfer_typer{   r"   r"   r#   serialize_http_requestj  s   z/BaseCRTRequestSerializer.serialize_http_requestN)r(   r)   r*   r   r"   r"   r"   r#   r   i  s    r   c                   @   sV   e Zd Z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S )BotocoreCRTRequestSerializerNc                 C   sp   || _ |du r	i }| || |jdi || _| jjjd| j | jjjd| j | jjjd| j	 dS )a  Serialize CRT HTTP request using botocore logic
        It also takes into account configuration from both the session
        and any keyword arguments that could be passed to
        `Session.create_client()` when serializing the request.

        :type session: botocore.session.Session

        :type client_kwargs: Optional[Dict[str, str]])
        :param client_kwargs: The kwargs for the botocore
            s3 client initialization.
        Nzrequest-created.s3.*zafter-call.s3.*zbefore-send.s3.*r"   )
Z_session_resolve_client_configZcreate_client_clientri   eventsregister_capture_http_request+_change_response_to_serialized_http_request_make_fake_http_response)r    sessionclient_kwargsr"   r"   r#   r$   z  s   


z%BotocoreCRTRequestSerializer.__init__c                 C   sP   d }|  r
|  }d|v r|d }ttd}|r||}||d< d|d< d S )Nconfig)Zsignature_versionZs3Zservice_name)Zget_default_client_configr   r   merge)r    r   r   Zuser_provided_configZclient_configr"   r"   r#   r     s   

z3BotocoreCRTRequestSerializer._resolve_client_configc           
      C   s   t |j}|j}|jr| d|j }g }|j D ]\}}t|tr+|||f q||t|df qt	j
|}d }|jrPt|jdrK|j}nt|j}t	j
j|j|||d}	|	S )N?zutf-8seek)methodpathheadersZbody_stream)r   urlr   queryr   items
isinstancestrrq   awscrthttpZHttpHeadersbodyhasattrr   ZHttpRequestr   )
r    Zaws_request	url_partsZcrt_pathZheaders_listnamevalueZcrt_headersZcrt_body_streamcrt_requestr"   r"   r#   _crt_request_from_aws_request  s,   


z:BotocoreCRTRequestSerializer._crt_request_from_aws_requestc                 C   sT   |  |}|jdd u rt|j}|jd|j |jdd ur(|jd |S )NhostzContent-MD5)r   r   getr   r   setnetlocremove)r    botocore_http_requestr   r   r"   r"   r#   _convert_to_crt_http_request  s   

z9BotocoreCRTRequestSerializer._convert_to_crt_http_requestc                 K   s   ||j d< d S )Nhttp_request)contextr    requestrf   r"   r"   r#   r     rV   z2BotocoreCRTRequestSerializer._capture_http_requestc                 K   s   |d }|  |d< d S )Nr   HTTPRequest)Zprepare)r    r   parsedrf   r   r"   r"   r#   r     s   zHBotocoreCRTRequestSerializer._change_response_to_serialized_http_requestc                 K   s   t jd di tdS )N       )botocoreZ
awsrequestZAWSResponseFakeRawResponser   r"   r"   r#   r     s   z5BotocoreCRTRequestSerializer._make_fake_http_responsec                 C   s&   t | j|d|j|jd|jd S )N)ZBucketKeyr   r"   )getattrr   rI   rJ   rL   )r    Zclient_methodrx   r"   r"   r#   _get_botocore_http_request  s   z7BotocoreCRTRequestSerializer._get_botocore_http_requestc                 C   s   |  ||jj}| |}|S r   )r   ri   rx   r   )r    r   r{   r   r   r"   r"   r#   r     s
   
z3BotocoreCRTRequestSerializer.serialize_http_requestr   )r(   r)   r*   r$   r   r   r   r   r   r   r   r   r"   r"   r"   r#   r   y  s    
r   c                   @   s   e Zd ZdddZdS )r      Nc                 c   s    	 |  |}|sd S |V  qr   )read)r    ZamtZdecode_contentchunkr"   r"   r#   stream  s   
zFakeRawResponse.stream)r   N)r(   r)   r*   r   r"   r"   r"   r#   r     s    r   c                   @   sd   e Zd ZdZdddZedd Zdd Zdd	d
ZdddZ	dd Z
dddZdd Zdd ZdS )rm   z-A helper class for managing CRTTransferFutureNc                 C   s0   || _ || _t | _d | _d | _t | _d S r   )	rh   _s3_requestr   r   r   
_exception_crt_futureEvent_done_event)r    rh   
s3_requestr"   r"   r#   r$     s   
zCRTTransferCoordinator.__init__c                 C   r   r   )r   r?   r"   r"   r#   r     r   z!CRTTransferCoordinator.s3_requestc                 C   rc   r   )r   r   r?   r"   r"   r#   set_done_callbacks_complete   rV   z2CRTTransferCoordinator.set_done_callbacks_completec                 C   r   r   )r   waitr   r"   r"   r#   r_     r   z<CRTTransferCoordinator.wait_until_on_done_callbacks_completeFc                 C   sN   | j  |  r
|r|| _W d    d S W d    d S 1 s w   Y  d S r   )r   rW   r   )r    r   r   r"   r"   r#   ru     s   "z$CRTTransferCoordinator.set_exceptionc                 C   s   | j r
| j   d S d S r   )r   rF   r?   r"   r"   r#   rF     s   zCRTTransferCoordinator.cancelc              	   C   sp   | j r| j z$z| j| W n ty   |    w W | jr#d | _| j| d S | jr1d | _| j| w r   )r   r   r\   ra   rF   r   r   r"   r"   r#   r\     s   zCRTTransferCoordinator.resultc                 C   s   | j d u rdS | j  S )NF)r   rW   r?   r"   r"   r#   rW     s   

zCRTTransferCoordinator.donec                 C   s   || _ | j j| _d S r   )r   Zfinished_futurer   )r    r   r"   r"   r#   rv   !  s   z%CRTTransferCoordinator.set_s3_requestr~   r   r   )r(   r)   r*   r   r$   r   r   r   r_   ru   rF   r\   rW   rv   r"   r"   r"   r#   rm     s    




rm   c                   @   s*   e Zd Zdd Zdd Z		dddZdS )	r7   c                 C   s   || _ || _d S r   )_request_serializer	_os_utils)r    r=   Zos_utilsr"   r"   r#   r$   '  s   
zS3ClientArgsCreator.__init__c              
   C   s   d }d }t t| tj}g }	|tjkr+|j}
| j|
}t||
|| j}|		| n|tj
kr>|j}| j|}||jd< | j||}||||| |d|	|| |ddS )NZContentLengthrW   progress)r   typerecv_filepathsend_filepathr}   Zon_progress)r   r   upperDEFAULTZ
GET_OBJECTrK   r   Zget_temp_filenameRenameTempFileHandlerrq   Z
PUT_OBJECTZget_file_sizerL   r   r   rs   )r    rw   rx   rY   r{   ry   r   r   Zs3_meta_request_typeZon_done_before_callsZfinal_filepathZfile_ondone_callZdata_lenr   r"   r"   r#   rt   +  s:   




z)S3ClientArgsCreator.get_make_request_argsNc                    s    fdd}|S )Nc                     sf   g }d ur
|7 }|t 7 } d ur| 7 }|D ]}dkr)|| d d q|| i | qd S )Nr   r   )Zbytes_transferred)r   )rE   rf   Zcallbacks_listcallbackrj   before_subscriberscallback_typer{   r"   r#   invoke_all_callbacksV  s   zBS3ClientArgsCreator.get_crt_callback.<locals>.invoke_all_callbacksr"   )r    r{   r   r   rj   r   r"   r   r#   rs   O  s   z$S3ClientArgsCreator.get_crt_callbackr~   )r(   r)   r*   r$   rt   rs   r"   r"   r"   r#   r7   &  s    (r7   c                   @      e Zd Zdd Zdd ZdS )r   c                 C   s   || _ || _|| _|| _d S r   )r   _final_filename_temp_filenamer5   )r    rY   Zfinal_filenameZtemp_filenamer>   r"   r"   r#   r$   j  s   
zRenameTempFileHandler.__init__c              
   K   sx   |d }|r| j | j d S z| j | j| j W d S  ty; } z| j | j | j| W Y d }~d S d }~ww )Nrk   )r5   Zremove_filer   Zrename_filer   rb   r   ru   )r    rf   rk   r|   r"   r"   r#   r'   p  s   
zRenameTempFileHandler.__call__Nr(   r)   r*   r$   r'   r"   r"   r"   r#   r   i  s    r   c                   @   r   )rp   c                 C   s
   || _ d S r   )r   rX   r"   r"   r#   r$     r   zAfterDoneHandler.__init__c                 K   rc   r   )r   r   re   r"   r"   r#   r'     rV   zAfterDoneHandler.__call__Nr   r"   r"   r"   r#   rp     s    rp   )8loggingr   ior   Zawscrt.httpr   Zbotocore.awsrequestr   Zbotocore.sessionZawscrt.authr   r   Z	awscrt.ior   r   r   r   r	   Z	awscrt.s3r
   r   r   r   Zbotocore.compatr   Zbotocore.configr   Zbotocore.exceptionsr   Zs3transfer.constantsr   r   Zs3transfer.exceptionsr   Zs3transfer.futuresr   r   Zs3transfer.utilsr   r   r   	getLoggerr(   loggerr   r2   r3   rn   ro   r   r   r   rm   r7   r   rp   r"   r"   r"   r#   <module>   sJ   


V &o	5C