U
    •uÄcº@  ã                   @   s,  d dl mZ d dlZe ¡  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 d d
lmZmZmZ d dlmZ zd dlmZ W n  ek
rÀ   d dlmZ Y nX d dlmZ G dd„ dejƒZG dd„ dejƒZG dd„ dejƒZG dd„ dejƒZ G dd„ dejƒZ!dS )é    )ÚMockN)Úwildcards2patterns)ÚWebSocketServerFactory)ÚWebSocketServerProtocol)ÚWebSocketClientProtocol)ÚTransportDetails©ÚConnectingRequest)ÚFailure)ÚConnectionDoneÚConnectionAbortedÚConnectionLost)Úunittest)ÚStringTransport)ÚFakeTransportc                   @   s@   e Zd ZdZdd„ Zdd„ Zdd„ Zdd	„ Zd
d„ Zdd„ Z	dS )ÚExceptionHandlingTestsz^
    Tests that we format various exception variations properly during
    connectionLost
    c                 C   s(   t ƒ | _tƒ | _| j| j_tƒ | j_d S ©N)r   Úfactoryr   Úprotor   Úlog©Úself© r   úqc:\users\psant\downloads\hotstarbackend\school\venv38\Lib\site-packages\autobahn\twisted\test\test_tx_protocol.pyÚsetUp7   s    
zExceptionHandlingTests.setUpc                 C   s6   | j j| j j| j j| j jfD ]}|d k	r| ¡  qd S r   ©r   ZautoPingPendingCallZautoPingTimeoutCallZopenHandshakeTimeoutCallZcloseHandshakeTimeoutCallÚcancel©r   Úcallr   r   r   ÚtearDown=   s    üzExceptionHandlingTests.tearDownc                 C   sH   | j  ¡  | j  ttƒ ƒ¡ d dd„ | j jjD ƒ¡}|  d|k¡ d S )Nú c                 S   s   g | ]}t |d  ƒ‘qS ©é   ©Ústr©Ú.0Úxr   r   r   Ú
<listcomp>M   s     z?ExceptionHandlingTests.test_connection_done.<locals>.<listcomp>zclosed cleanly)	r   Ú_connectionMadeÚconnectionLostr
   r   Újoinr   Ú
mock_callsÚ
assertTrue©r   Úmessagesr   r   r   Útest_connection_doneG   s    
z+ExceptionHandlingTests.test_connection_donec                 C   sH   | j  ¡  | j  ttƒ ƒ¡ d dd„ | j jjD ƒ¡}|  d|k¡ d S )Nr    c                 S   s   g | ]}t |d  ƒ‘qS r!   r#   r%   r   r   r   r(   V   s     zBExceptionHandlingTests.test_connection_aborted.<locals>.<listcomp>z	 aborted )	r   r)   r*   r
   r   r+   r   r,   r-   r.   r   r   r   Útest_connection_abortedP   s    
z.ExceptionHandlingTests.test_connection_abortedc                 C   sH   | j  ¡  | j  ttƒ ƒ¡ d dd„ | j jjD ƒ¡}|  d|k¡ d S )Nr    c                 S   s   g | ]}t |d  ƒ‘qS r!   r#   r%   r   r   r   r(   _   s     z?ExceptionHandlingTests.test_connection_lost.<locals>.<listcomp>ú
 was lost ©	r   r)   r*   r
   r   r+   r   r,   r-   r.   r   r   r   Útest_connection_lostY   s    
z+ExceptionHandlingTests.test_connection_lostc                 C   sX   | j  ¡  | j  ttdƒƒ¡ d dd„ | j jjD ƒ¡}|  d|k¡ |  d|k¡ d S )NZ	greetingsr    c                 S   s$   g | ]}t |d  ƒt |d ƒ ‘qS )r"   é   r#   r%   r   r   r   r(   h   s     zCExceptionHandlingTests.test_connection_lost_arg.<locals>.<listcomp>r2   r3   r.   r   r   r   Útest_connection_lost_argb   s
    
z/ExceptionHandlingTests.test_connection_lost_argN)
Ú__name__Ú
__module__Ú__qualname__Ú__doc__r   r   r0   r1   r4   r6   r   r   r   r   r   1   s   
			r   c                   @   s   e Zd ZdZdd„ ZdS )ÚHixie76RejectionTestsz@
    Hixie-76 should not be accepted by an Autobahn server.
    c                 C   s^   t ƒ }tƒ }tƒ }||_||_d}d|_| ¡  ||_| ¡  |  	d|j
¡ |  	d|j
¡ dS )zO
        A handshake from a client only supporting Hixie-76 will fail.
        sè   GET /demo HTTP/1.1
Host: example.com
Connection: Upgrade
Sec-WebSocket-Key2: 12998 5 Y3 1  .P00
Sec-WebSocket-Protocol: sample
Upgrade: WebSocket
Sec-WebSocket-Key1: 4 @1  46546xW%0l 1 5
Origin: http://example.com

^n:ds[4Ur   s   HTTP/1.1 400s   Hixie76 protocol not supportedN)r   r   r   r   Ú	transportÚopenHandshakeTimeoutr)   ÚdataZprocessHandshakeÚassertInZ_written)r   ÚtÚfÚpZhttp_requestr   r   r   Útest_handshake_failsq   s    z*Hixie76RejectionTests.test_handshake_failsN)r7   r8   r9   r:   rC   r   r   r   r   r;   m   s   r;   c                   @   sp   e Zd 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d„ Zdd„ Zdd„ ZdS )ÚWebSocketOriginMatchingzD
    Test that we match Origin: headers properly, when asked to
    c                 C   sN   t ƒ | _| jjddgd� tƒ | _tƒ | j_| j| j_tƒ | j_| j 	¡  d S )Nz127.0.0.1:*z*.example.com:*©ÚallowedOrigins©
r   r   ÚsetProtocolOptionsr   r   r   r<   r   ÚfailHandshaker)   r   r   r   r   r   ‹   s    ÿ


zWebSocketOriginMatching.setUpc                 C   s6   | j j| j j| j j| j jfD ]}|d k	r| ¡  qd S r   r   r   r   r   r   r   –   s    üz WebSocketOriginMatching.tearDownc                 C   sd   d  dddddddd	dg	¡| j_| j ¡  |  | jjjd
¡ | jjjd d d }|  d|k¡ d S )Nó   
ó   GET /ws HTTP/1.1ó   Host: www.example.comó   Sec-WebSocket-Version: 13s,   Origin: http://www.example.com.malicious.comó,   Sec-WebSocket-Extensions: permessage-deflateó+   Sec-WebSocket-Key: tXAxWFUqnhi86Ajj7dRY5g==ó   Connection: keep-alive, Upgradeó   Upgrade: websocketúHandshake should have failedr   r"   únot allowed)r+   r   r>   ÚconsumeDatar-   rI   Úcalledr,   ©r   Úargr   r   r   Útest_match_full_origin    s    ÷

z.WebSocketOriginMatching.test_match_full_originc                 C   s”   | j jdgd� | j j| j_| j j| j_d| j _d ddddd	d
dddg	¡| j_| j ¡  |  	| jj
jd¡ | jj
jd d d }|  	d|k¡ d S )Nzhttp://*.example.com:*rE   FrJ   rK   rL   rM   ó   Origin: https://www.example.comrN   rO   rP   rQ   rR   r   r"   rS   )r   rH   ZallowedOriginsPatternsr   rF   ÚisSecurer+   r>   rT   r-   rI   rU   r,   rV   r   r   r   Útest_match_wrong_scheme_origin²   s*    ÿ÷

z6WebSocketOriginMatching.test_match_wrong_scheme_originc                 C   sP   d| j _d| j _d dddddd	d
ddg	¡| j_| j ¡  |  | jjj	d¡ d S )NTé»  rJ   rK   rL   rM   rY   rN   rO   rP   rQ   úHandshake should have succeeded)
r   rZ   Úportr+   r   r>   rT   ÚassertFalserI   rU   r   r   r   r   Útest_match_origin_secure_schemeÎ   s    ÷

z7WebSocketOriginMatching.test_match_origin_secure_schemec                 C   s`   | j jdgd� d| j _d| j _d dddd	d
ddddg	¡| j_| j ¡  |  | jj	j
d¡ dS )z1
        Test the examples from the docs
        z*://*.example.com:*rE   Tr\   rJ   rK   rL   rM   ó   Origin: http://www.example.comrN   rO   rP   rQ   r]   N)r   rH   rZ   r^   r+   r   r>   rT   r_   rI   rU   r   r   r   r   Ú'test_match_origin_documentation_exampleà   s$    ÿ÷

z?WebSocketOriginMatching.test_match_origin_documentation_examplec                 C   sB   ddl m}m} tdgƒ}dD ]}|  |||ƒdd|ƒ|¡ qdS )z>
        All the example origins from RFC6454 (3.2.1)
        r   ©Ú_is_same_originÚ_url_to_originz*example.com:*)zhttp://example.com/zhttp://example.com:80/zhttp://example.com/path/filezhttp://example.com/;semi=truez//example.com/zhttp://@example.comÚhttpéP   N)Úautobahn.websocket.protocolrd   re   r   r-   ©r   rd   re   ÚpolicyÚurlr   r   r   Útest_match_origin_examplesø   s    
z2WebSocketOriginMatching.test_match_origin_examplesc                 C   s@   ddl m}m} tdgƒ}dD ]}|  |||ƒdd|ƒ¡ qdS )zI
        All the example 'not-same' origins from RFC6454 (3.2.1)
        r   rc   úexample.com)zhttp://ietf.org/zhttp://example.org/zhttps://example.com/zhttp://example.com:8080/zhttp://www.example.com/rf   rg   N)rh   rd   re   r   r_   ri   r   r   r   Ú"test_match_origin_counter_examples	  s    
z:WebSocketOriginMatching.test_match_origin_counter_examplesc                 C   sf   ddl m}m} tdgƒ}|  ||dƒdd|ƒ¡ |  ||dƒdd|ƒ¡ |  ||d	ƒdd|ƒ¡ d S )
Nr   rc   zhttp://*example.com:80zhttp://example.com:80rf   rg   zhttp://example.com:81éQ   zhttps://example.com:80)rh   rd   re   r   r-   r_   )r   rd   re   rj   r   r   r   Útest_match_origin_edge  s    
ÿÿÿz.WebSocketOriginMatching.test_match_origin_edgec                 C   s0   ddl m} |  |dƒd¡ |  |dƒd¡ d S )Nr   ©re   zhttp://example.com)rf   rm   rg   zhTTp://example.com©rh   re   ÚassertEqual©r   re   r   r   r   Útest_origin_from_url%  s    þþz,WebSocketOriginMatching.test_origin_from_urlc                 C   s    ddl m} |  d|dƒ¡ d S )Nr   rq   Únullzfile:///etc/passwdrr   rt   r   r   r   Útest_origin_file3  s    z(WebSocketOriginMatching.test_origin_filec                 C   s„   ddl m}m} |  d|dƒ¡ |  ||dƒddg ƒ¡ |  ||dƒddg ƒ¡ |  ||dƒddg ƒ¡ |  ||dƒd dg ƒ¡ d S )Nr   rc   rv   rf   rg   ÚhttpsÚ )rh   rd   re   rs   r_   )r   rd   re   r   r   r   Útest_origin_null7  s    ÿÿÿÿz(WebSocketOriginMatching.test_origin_nullN)r7   r8   r9   r:   r   r   rX   r[   r`   rb   rl   rn   rp   ru   rw   rz   r   r   r   r   rD   †   s   
rD   c                   @   s(   e Zd ZdZdd„ Zdd„ Zdd„ ZdS )	ÚWebSocketXForwardedForzR
    Test that (only) a trusted X-Forwarded-For can replace the peer address.
    c                 C   sJ   t ƒ | _| jjdd� tƒ | _tƒ | j_| j| j_tƒ | j_| j 	¡  d S )Nr5   )ZtrustXForwardedForrG   r   r   r   r   r   M  s    ÿ


zWebSocketXForwardedFor.setUpc                 C   s6   | j j| j j| j j| j jfD ]}|d k	r| ¡  qd S r   r   r   r   r   r   r   X  s    üzWebSocketXForwardedFor.tearDownc                 C   sB   d  dddddddd	d
dg
¡| j_| j ¡  |  | jjdd¡ d S )NrJ   rK   rL   ra   rM   rN   rO   rP   rQ   s0   X-Forwarded-For: 1.2.3.4, 2.3.4.5, 111.222.33.44z2.3.4.5zQThe second address in X-Forwarded-For should have been picked as the peer address)r+   r   r>   rT   ÚassertEqualsÚpeerr   r   r   r   Útest_trusted_addressesb  s$    ö

 þz-WebSocketXForwardedFor.test_trusted_addressesN)r7   r8   r9   r:   r   r   r~   r   r   r   r   r{   H  s   
r{   c                   @   s0   e Zd ZdZdd„ Zdd„ Zdd„ Zdd	„ Zd
S )ÚOnConnectingTestsz¾
    Tests related to onConnecting callback

    These tests are testing generic behavior, but are somewhat tied to
    'a framework' so we're just testing using Twisted-specifics here.
    c                    s�   d‰ G ‡ fdd„dt ƒ}|ƒ }tƒ |_| ¡ }|  |¡ |  t|jjƒdk¡ d}t	t|jjƒƒD ]}ˆ t
|jj| ƒkrbd}qb|  |d¡ d S )Nz	bad stuffc                       s&   e Zd ZdZdZeƒ Z‡ fdd„ZdS )zDOnConnectingTests.test_on_connecting_client_fails.<locals>.TestProtoNTc                    s   t ˆ ƒ‚d S r   )ÚRuntimeError©r   Ztransport_details©Z	MAGIC_STRr   r   ÚonConnecting‡  s    zQOnConnectingTests.test_on_connecting_client_fails.<locals>.TestProto.onConnecting)r7   r8   r9   ÚstateÚwasCleanr   r   rƒ   r   r‚   r   r   Ú	TestProto‚  s   r†   r   FTz!MAGIC_STR not found when expected)r   r   r<   ÚstartHandshakeÚsuccessResultOfr-   Úlenr   r,   Úranger$   )r   r†   r   ÚdZmagic_foundÚir   r‚   r   Útest_on_connecting_client_fails~  s    
z1OnConnectingTests.test_on_connecting_client_failsc                 C   sn   G dd„ dt ƒ}|ƒ }tƒ |_tƒ |_| ¡  | ¡ }|  |¡}|  d|j	¡ |  d|j
¡ |  d|j¡ d S )Nc                   @   s.   e Zd ZdZdZg ZdZdZeƒ Z	dd„ Z
dS )zFOnConnectingTests.test_on_connecting_client_success.<locals>.TestProtoNTé   é   c                 S   s   t dddd�S )Nrm   r\   ú/ws©Úhostr^   Úresourcer   r�   r   r   r   rƒ      s
    ýzSOnConnectingTests.test_on_connecting_client_success.<locals>.TestProto.onConnecting)r7   r8   r9   r„   r…   ZperMessageCompressionOffersÚversionr=   r   r   rƒ   r   r   r   r   r†   ˜  s   r†   rm   r\   r�   )r   r   r<   r   r   r)   r‡   rˆ   rs   r’   r^   r“   )r   r†   r   r‹   Úreqr   r   r   Ú!test_on_connecting_client_success–  s    
z3OnConnectingTests.test_on_connecting_client_successc                 C   s   t t jddi d�}t|ƒ d S )Nrm   F)Zchannel_typer}   Z	is_secureZ
channel_id)r   ZCHANNEL_TYPE_FUNCTIONr$   )r   Údetailsr   r   r   Útest_str_transport²  s    üz$OnConnectingTests.test_str_transportc                 C   s   t dddd�}t|ƒ d S )Nrm   Z1234r�   r‘   )r	   r$   )r   r•   r   r   r   Útest_str_connecting¼  s    z%OnConnectingTests.test_str_connectingN)r7   r8   r9   r:   r�   r–   r˜   r™   r   r   r   r   r   v  s
   
r   )"Zunittest.mockr   ZtxaioZuse_twistedZautobahn.utilr   Zautobahn.twisted.websocketr   r   r   Zautobahn.wamp.typesr   Zautobahn.websocket.typesr	   Ztwisted.python.failurer
   Ztwisted.internet.errorr   r   r   Ztwisted.trialr   Ztwisted.internet.testingr   ÚImportErrorZtwisted.test.proto_helpersZautobahn.testutilr   ÚTestCaser   r;   rD   r{   r   r   r   r   r   Ú<module>   s,   < C.