o
    èó·iº@  ã                   @   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y_   d dlmZ Y nw 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   úlC:\Users\psant\Downloads\PAF backend\school\venv\Lib\site-packages\autobahn/twisted/test/test_tx_protocol.pyÚsetUp7   s   
zExceptionHandlingTests.setUpc                 C   ó6   | j j| j j| j j| j jfD ]
}|d ur| ¡  qd S r   ©r   ZautoPingPendingCallZautoPingTimeoutCallZopenHandshakeTimeoutCallZcloseHandshakeTimeoutCallÚcancel©r   Úcallr   r   r   ÚtearDown=   ó   ü€ùzExceptionHandlingTests.tearDownc                 C   óH   | j  ¡  | j  ttƒ ƒ¡ d dd„ | j jjD ƒ¡}|  d|v ¡ d S )Nú c                 S   ó   g | ]}t |d  ƒ‘qS ©é   ©Ústr©Ú.0Úxr   r   r   Ú
<listcomp>M   ó    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   ó   
z+ExceptionHandlingTests.test_connection_donec                 C   r"   )Nr#   c                 S   r$   r%   r'   r)   r   r   r   r,   V   r-   zBExceptionHandlingTests.test_connection_aborted.<locals>.<listcomp>z	 aborted )	r   r.   r/   r
   r   r0   r   r1   r2   r3   r   r   r   Útest_connection_abortedP   r6   z.ExceptionHandlingTests.test_connection_abortedc                 C   r"   )Nr#   c                 S   r$   r%   r'   r)   r   r   r   r,   _   r-   z?ExceptionHandlingTests.test_connection_lost.<locals>.<listcomp>ú
 was lost ©	r   r.   r/   r
   r   r0   r   r1   r2   r3   r   r   r   Útest_connection_lostY   r6   z+ExceptionHandlingTests.test_connection_lostc                 C   sX   | j  ¡  | j  ttdƒƒ¡ d dd„ | j jjD ƒ¡}|  d|v ¡ |  d|v ¡ 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>r8   r9   r3   r   r   r   Útest_connection_lost_argb   s
   
z/ExceptionHandlingTests.test_connection_lost_argN)
Ú__name__Ú
__module__Ú__qualname__Ú__doc__r   r    r5   r7   r:   r<   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processHandshakeZassertInZ_written)r   ÚtÚfÚpÚhttp_requestr   r   r   Útest_handshake_failsq   s   z*Hixie76RejectionTests.test_handshake_failsN)r=   r>   r?   r@   rI   r   r   r   r   rA   m   s    rA   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   rB   r   ÚfailHandshaker.   r   r   r   r   r   ‹   s   ÿ


zWebSocketOriginMatching.setUpc                 C   r   r   r   r   r   r   r   r    –   r!   z WebSocketOriginMatching.tearDownc                 C   sV   d  g d¢¡| j_| j ¡  |  | jjjd¡ | jjjd d d }|  d|v ¡ 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: websocketrP   úHandshake should have failedr   r&   únot allowed)r0   r   rD   ÚconsumeDatar2   rO   Úcalledr1   ©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 g d¢¡| j_| j ¡  |  	| jj
jd¡ | jj
jd d d }|  	d	|v ¡ d S )
Nzhttp://*.example.com:*rK   FrP   ©	rQ   rR   rS   s   Origin: https://www.example.comrT   rU   rV   rW   rP   rX   r   r&   rY   )r   rN   ZallowedOriginsPatternsr   rL   ÚisSecurer0   rD   rZ   r2   rO   r[   r1   r\   r   r   r   Útest_match_wrong_scheme_origin²   s   ÿ
z6WebSocketOriginMatching.test_match_wrong_scheme_originc                 C   sB   d| j _d| j _d g d¢¡| j_| j ¡  |  | jjj	d¡ d S )NTé»  rP   r_   úHandshake should have succeeded)
r   r`   Úportr0   r   rD   rZ   ÚassertFalserO   r[   r   r   r   r   Útest_match_origin_secure_schemeÎ   s
   
z7WebSocketOriginMatching.test_match_origin_secure_schemec                 C   sR   | j jdgd� d| j _d| j _d g d¢¡| j_| j ¡  |  | jj	j
d¡ dS )	z1
        Test the examples from the docs
        z*://*.example.com:*rK   Trb   rP   )	rQ   rR   rS   ó   Origin: http://www.example.comrT   rU   rV   rW   rP   rc   N)r   rN   r`   rd   r0   r   rD   rZ   re   rO   r[   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.protocolrj   rk   r   r2   ©r   rj   rk   Ú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   ri   úexample.com)zhttp://ietf.org/zhttp://example.org/zhttps://example.com/zhttp://example.com:8080/zhttp://www.example.com/rl   rm   N)rn   rj   rk   r   re   ro   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   ri   zhttp://*example.com:80zhttp://example.com:80rl   rm   zhttp://example.com:81éQ   zhttps://example.com:80)rn   rj   rk   r   r2   re   )r   rj   rk   rp   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   ©rk   zhttp://example.com)rl   rs   rm   zhTTp://example.com©rn   rk   ÚassertEqual©r   rk   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   rw   Únullzfile:///etc/passwdrx   rz   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   ri   r|   rl   rm   ÚhttpsÚ )rn   rj   rk   ry   re   )r   rj   rk   r   r   r   Útest_origin_null7  s   ÿÿÿÿz(WebSocketOriginMatching.test_origin_nullN)r=   r>   r?   r@   r   r    r^   ra   rf   rh   rr   rt   rv   r{   r}   r€   r   r   r   r   rJ   †   s    
rJ   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 )Nr;   )ZtrustXForwardedForrM   r   r   r   r   r   M  s   ÿ


zWebSocketXForwardedFor.setUpc                 C   r   r   r   r   r   r   r   r    X  r!   zWebSocketXForwardedFor.tearDownc                 C   s2   d  g d¢¡| j_| j ¡  |  | jjdd¡ d S )NrP   )
rQ   rR   rg   rS   rT   rU   rV   rW   s0   X-Forwarded-For: 1.2.3.4, 2.3.4.5, 111.222.33.44rP   z2.3.4.5zQThe second address in X-Forwarded-For should have been picked as the peer address)r0   r   rD   rZ   ZassertEqualsÚpeerr   r   r   r   Útest_trusted_addressesb  s   
þz-WebSocketXForwardedFor.test_trusted_addressesN)r=   r>   r?   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| ƒv r?d}q1|  |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 ˆ ƒ‚r   )ÚRuntimeError©r   Ztransport_details©Z	MAGIC_STRr   r   ÚonConnecting‡  s   zQOnConnectingTests.test_on_connecting_client_fails.<locals>.TestProto.onConnecting)r=   r>   r?   ÚstateÚwasCleanr   r   rˆ   r   r‡   r   r   Ú	TestProto‚  s
    r‹   r   FTz!MAGIC_STR not found when expected)r   r   rB   ÚstartHandshakeÚsuccessResultOfr2   Úlenr   r1   Ú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 )Nrs   rb   ú/ws©Úhostrd   Úresourcer   r†   r   r   r   rˆ      s
   ýzSOnConnectingTests.test_on_connecting_client_success.<locals>.TestProto.onConnecting)r=   r>   r?   r‰   rŠ   ZperMessageCompressionOffersÚversionrC   r   r   rˆ   r   r   r   r   r‹   ˜  s    r‹   rs   rb   r•   )r   r   rB   r   r   r.   rŒ   r�   ry   r—   rd   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 )Nrs   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 )Nrs   Z1234r•   r–   )r	   r(   )r   rš   r   r   r   Útest_str_connecting¼  s   z%OnConnectingTests.test_str_connectingN)r=   r>   r?   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   ZTestCaser   rA   rJ   r�   r„   r   r   r   r   Ú<module>   s0   ÿ< C.