U
    O8cV                     @   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"   2/tmp/pip-unpacked-wheel-6hpttf6a/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    credentialsr"   r"   r#   __call__0   s      z%CRTCredentialProviderAdapter.__call__c              
   C   sL   | j < | jd kr0| j }|d kr*t || _| jW  5 Q R  S Q R X 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k	rht }|rR|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_throughputr0   Zuse_sslverifyZevent_loop_groupZhost_resolverr.   Zproviderr2   r1   Ztls_ctx_optionsZclient_tls_optionZcredentails_provider_adapterZtarget_gbpsr"   r"   r#   create_s3_crt_clientB   sB    /
r4   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k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 kri }|d kri }t |||||d}| d|S )Nbucketkeyfileobj
extra_argssubscribersZ
get_objectr   _submit_transfer)r    rK   rL   rM   rN   rO   callargsr"   r"   r#   download   s    zCRTTransferManager.downloadc                 C   s6   |d kri }|d kri }t |||||d}| d|S )NrJ   Z
put_objectrP   )r    rM   rK   rL   rN   rO   rR   r"   r"   r#   upload   s    zCRTTransferManager.uploadc                 C   s4   |d kri }|d kri }t ||||d}| d|S )N)rK   rL   rN   rO   Zdelete_objectrP   )r    rK   rL   rN   rO   rR   r"   r"   r#   delete   s    zCRTTransferManager.deleteFc                 C   s   |  | d S r   rC   r    rH   r"   r"   r#   shutdown   s    zCRTTransferManager.shutdownc                 C   s    | j D ]}| s|  qd S r   )r;   donerH   r    coordinatorr"   r"   r#   _cancel_transfers   s    
z$CRTTransferManager._cancel_transfersc                 C   s   | j D ]}|  qd S r   )r;   resultrY   r"   r"   r#   _finish_transfers   s    
z$CRTTransferManager._finish_transfersc                 C   s   | j D ]}|  qd S r   )r;   %wait_until_on_done_callbacks_completerY   r"   r"   r#   _wait_transfers_done   s    
z'CRTTransferManager._wait_transfers_donec                 C   s\   |r|    z@z|   W n. tk
r6   |    Y n tk
rH   Y nX W 5 |   X d S r   )r[   r_   r]   KeyboardInterrupt	ExceptionrV   r"   r"   r#   rD      s    
zCRTTransferManager._shutdownc                 K   s   | j   d S r   )r=   releaser    kwargsr"   r"   r#   _release_semaphore  s    z%CRTTransferManager._release_semaphorec              
   C   s   | j g}t| jd}t| j||d}tf |}t|}|| zD| j  | j	
|d}|  | j	|||||}	| jjf |	}
W nJ tk
r } z,||d | j	j
|d|d}||d W 5 d }~X Y nX ||
 | j| |  jd7  _|S )	N)transfer_id)metarZ   ZqueuedTrX   )after_subscribers)error   )re   CRTTransferCoordinatorr>   CRTTransferMetaCRTTransferFutureAfterDoneHandlerappendr=   acquirer:   get_crt_callbackget_make_request_argsr8   Zmake_requestra   set_exceptionset_s3_requestr;   )r    request_type	call_argson_done_after_callsrZ   
componentsfutureZ	afterdoneZ	on_queuedZcrt_callargsZcrt_s3_requesteon_doner"   r"   r#   rQ     sF    



   
z#CRTTransferManager._submit_transfer)N)NN)NN)NN)F)F)r)   r*   r+   r$   rB   rI   rS   rT   rU   rW   r[   r]   r_   rD   re   rQ   r"   r"   r"   r#   r5      s   
   




r5   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 )rl   z*Holds metadata about the CRTTransferFutureNc                 C   s   || _ || _i | _d S r   )_transfer_id
_call_args_user_context)r    rf   rv   r"   r"   r#   r$   1  s    zCRTTransferMeta.__init__c                 C   s   | j S r   )r}   rA   r"   r"   r#   rv   6  s    zCRTTransferMeta.call_argsc                 C   s   | j S r   )r|   rA   r"   r"   r#   rf   :  s    zCRTTransferMeta.transfer_idc                 C   s   | j S r   )r~   rA   r"   r"   r#   user_context>  s    zCRTTransferMeta.user_context)NN)	r)   r*   r+   __doc__r$   propertyrv   rf   r   r"   r"   r"   r#   rl   .  s   


rl   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 )rm   Nc                 C   s    || _ |dk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)_metarl   _coordinator)r    rg   rZ   r"   r"   r#   r$   D  s    	zCRTTransferFuture.__init__c                 C   s   | j S r   )r   rA   r"   r"   r#   rg   R  s    zCRTTransferFuture.metac                 C   s
   | j  S r   )r   rX   rA   r"   r"   r#   rX   V  s    zCRTTransferFuture.donec                 C   s   | j | d S r   )r   r\   r    timeoutr"   r"   r#   r\   Y  s    zCRTTransferFuture.resultc                 C   s   | j   d S r   )r   rH   rA   r"   r"   r#   rH   \  s    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)rX   r   r   rs   )r    	exceptionr"   r"   r#   rs   _  s
    zCRTTransferFuture.set_exception)NN)N)
r)   r*   r+   r$   r   rg   rX   r\   rH   rs   r"   r"   r"   r#   rm   C  s   


rm   c                   @   s   e Zd Zdd ZdS )BaseCRTRequestSerializerc                 C   s   t ddS )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()N)NotImplementedError)r    transfer_typery   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   sl   || _ |dkri }| || |jf || _| 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.*)
_session_resolve_client_configZcreate_client_clientrg   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|k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 ]6\}}t|trV|||f q4||t|df q4t	j
|}d }|jrt|jdr|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strro   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 kr4t|j}|jd|j |jdd k	rP|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    requestrd   r"   r"   r#   r     s    z2BotocoreCRTRequestSerializer._capture_http_requestc                 K   s   |d }|  |d< d S )Nr   HTTPRequest)prepare)r    r   parsedrd   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|f |j|jd|jd S )N)ZBucketZKeyr   )getattrr   rK   rL   rN   )r    Zclient_methodrv   r"   r"   r#   _get_botocore_http_request  s     z7BotocoreCRTRequestSerializer._get_botocore_http_requestc                 C   s   |  ||jj}| |}|S r   )r   rg   rv   r   )r    r   ry   r   r   r"   r"   r#   r     s     
z3BotocoreCRTRequestSerializer.serialize_http_request)N)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q|V  q d S r   )read)r    amt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 )rk   z-A helper class for managing CRTTransferFutureNc                 C   s0   || _ || _t | _d | _d | _t | _d S r   )	rf   _s3_requestr   r   r   
_exception_crt_futureEvent_done_event)r    rf   
s3_requestr"   r"   r#   r$     s    
zCRTTransferCoordinator.__init__c                 C   s   | j S r   )r   rA   r"   r"   r#   r     s    z!CRTTransferCoordinator.s3_requestc                 C   s   | j   d S r   )r   r   rA   r"   r"   r#   set_done_callbacks_complete   s    z2CRTTransferCoordinator.set_done_callbacks_completec                 C   s   | j | d S r   )r   waitr   r"   r"   r#   r^     s    z<CRTTransferCoordinator.wait_until_on_done_callbacks_completeFc              	   C   s(   | j  |  r|r|| _W 5 Q R X d S r   )r   rX   r   )r    r   r   r"   r"   r#   rs     s    z$CRTTransferCoordinator.set_exceptionc                 C   s   | j r| j   d S r   )r   rH   rA   r"   r"   r#   rH     s    zCRTTransferCoordinator.cancelc              	   C   s`   | j r| j z4z| j| W n tk
r<   |    Y nX W 5 | jrNd | _| j| X d S r   )r   r   r   r\   r`   rH   r   r"   r"   r#   r\     s    zCRTTransferCoordinator.resultc                 C   s   | j d krdS | j  S )NF)r   rX   rA   r"   r"   r#   rX     s    
zCRTTransferCoordinator.donec                 C   s   || _ | j j| _d S r   )r   Zfinished_futurer   )r    r   r"   r"   r#   rt   !  s    z%CRTTransferCoordinator.set_s3_request)NN)N)F)N)r)   r*   r+   r   r$   r   r   r   r^   rs   rH   r\   rX   rt   r"   r"   r"   r#   rk     s   




rk   c                   @   s&   e Zd Zdd Zdd ZdddZdS )	r9   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rV|j}
| j|
}t||
|| j}|		| n&|tj
kr||j}| j|}||jd< | j||}||||| |d|	|| |ddS )NZContentLengthrX   progress)r   typerecv_filepathsend_filepathr{   Zon_progress)r   r   upperDEFAULTZ
GET_OBJECTrM   r   Zget_temp_filenameRenameTempFileHandlerro   Z
PUT_OBJECTZget_file_sizerN   r   r   rq   )r    ru   rv   rZ   ry   rw   r   r   Zs3_meta_request_typeZon_done_before_callsZfinal_filepathZfile_ondone_callZdata_lenr   r"   r"   r#   rr   +  sL      
   

    
z)S3ClientArgsCreator.get_make_request_argsNc                    s    fdd}|S )Nc                     sb   g }d k	r|7 }|t 7 } d k	r2| 7 }|D ]&}dkrR|| d d q6|| | q6d S )Nr   r   )Zbytes_transferred)r   )rG   rd   Zcallbacks_listcallbackrh   before_subscriberscallback_typery   r"   r#   invoke_all_callbacksV  s    zBS3ClientArgsCreator.get_crt_callback.<locals>.invoke_all_callbacksr"   )r    ry   r   r   rh   r   r"   r   r#   rq   O  s    z$S3ClientArgsCreator.get_crt_callback)NN)r)   r*   r+   r$   rr   rq   r"   r"   r"   r#   r9   &  s
   (  r9   c                   @   s   e Zd Zdd Zdd ZdS )r   c                 C   s   || _ || _|| _|| _d S r   )r   _final_filename_temp_filenamer7   )r    rZ   Zfinal_filenameZtemp_filenamer@   r"   r"   r#   r$   j  s    zRenameTempFileHandler.__init__c              
   K   st   |d }|r| j | j nTz| j | j| j W n< tk
rn } z| j | j | j| W 5 d }~X Y nX d S )Nri   )r7   Zremove_filer   Zrename_filer   ra   r   rs   )r    rd   ri   rz   r"   r"   r#   r(   p  s     zRenameTempFileHandler.__call__Nr)   r*   r+   r$   r(   r"   r"   r"   r#   r   i  s   r   c                   @   s   e Zd Zdd Zdd ZdS )rn   c                 C   s
   || _ d S r   )r   rY   r"   r"   r#   r$     s    zAfterDoneHandler.__init__c                 K   s   | j   d S r   )r   r   rc   r"   r"   r#   r(     s    zAfterDoneHandler.__call__Nr   r"   r"   r"   r#   rn     s   rn   )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   r4   r5   rl   rm   r   r   r   rk   r9   r   rn   r"   r"   r"   r#   <module>   sH   


V &o	5C