o
    ¯bXD  ã                   @   sº   d Z ddlZddlZddl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 G d	d
„ d
ejƒZG dd„ dƒZG dd„ de	jƒZG dd„ dejƒZG dd„ dejƒZdS )z_
Tests for L{twisted.protocol.socks}, an implementation of the SOCKSv4 and
SOCKSv4a protocols.
é    N)ÚaddressÚdefer)ÚDNSLookupError)Úsocks)Ú	iterbytes)Úproto_helpers)Úunittestc                   @   s,   e Zd ZdZdZdd„ Zdd„ Zdd„ ZdS )	ÚStringTCPTransportFNc                 C   s   | j S )N)Úpeer©Úself© r   ú9/usr/lib/python3/dist-packages/twisted/test/test_socks.pyÚgetPeer   s   zStringTCPTransport.getPeerc                 C   s   t  ddd¡S )NÚTCPz2.3.4.5é*   )r   ÚIPv4Addressr   r   r   r   ÚgetHost   s   zStringTCPTransport.getHostc                 C   s
   d| _ d S )NT)ÚstringTCPTransport_closingr   r   r   r   ÚloseConnection   s   
z!StringTCPTransport.loseConnection)Ú__name__Ú
__module__Ú__qualname__r   r
   r   r   r   r   r   r   r   r	      s    r	   c                   @   s    e Zd ZdZdd„ Zdd„ ZdS )ÚFakeResolverReactorzP
    Bare-bones reactor with deterministic behavior for the resolve method.
    c                 C   s
   || _ dS )zÃ
        @type names: L{dict} containing L{str} keys and L{str} values.
        @param names: A hostname to IP address mapping. The IP addresses are
            stringified dotted quads.
        N)Únames)r   r   r   r   r   Ú__init__'   s   
zFakeResolverReactor.__init__c              
   C   s>   z	t  | j| ¡W S  ty   t  td| d¡ ƒ¡ Y S w )zQ
        Resolve a hostname by looking it up in the C{names} dictionary.
        z"FakeResolverReactor couldn't find zutf-8)r   Úsucceedr   ÚKeyErrorÚfailr   Údecode)r   Úhostnamer   r   r   Úresolve/   s   ÿÿÿzFakeResolverReactor.resolveN)r   r   r   Ú__doc__r   r!   r   r   r   r   r   "   s    r   c                   @   s$   e Zd ZdZdZdd„ Zdd„ ZdS )ÚSOCKSv4DriverNc                 G   s:   ||Ž }t ƒ |_t d||¡|j_| ¡  || _t |¡S )Nr   )	r	   Ú	transportr   r   r
   ÚconnectionMadeÚdriver_outgoingr   r   )r   ÚhostÚportÚklassÚargsÚprotor   r   r   ÚconnectClassD   s   
zSOCKSv4Driver.connectClassc                 G   s(   ||Ž }|| _ |dkrd}t d|f¡S )Nr   éÒ  ú6.7.8.9)Údriver_listenr   r   )r   r(   r)   r*   Úfactoryr   r   r   ÚlistenClassM   s
   zSOCKSv4Driver.listenClass)r   r   r   r&   r/   r,   r1   r   r   r   r   r#   =   s
    	r#   c                   @   óP   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S )ÚConnectTestszV
    Tests for SOCKS and SOCKSv4a connect requests using the L{SOCKSv4} protocol.
    c                 C   ó0   t ƒ | _tƒ | j_| j ¡  tddiƒ| j_d S ©Ns	   localhostú	127.0.0.1©r#   Úsockr	   r$   r%   r   Úreactorr   r   r   r   ÚsetUp[   ó   

zConnectTests.setUpc                 C   s(   | j j}|d ur|  |jjd¡ d S d S )Nz-Outgoing SOCKS connections need to be closed.)r8   r&   Ú
assertTruer$   r   )r   Úoutgoingr   r   r   ÚtearDowna   s   þÿzConnectTests.tearDownc              	   C   sÖ   | j  t dddd¡t d¡ d d ¡ | j j ¡ }| j j ¡  |  	|t ddd	d¡t d¡ ¡ |  
| j jj¡ |  | j j¡ | j  d
¡ |  	| j jj ¡ d
¡ | j j d¡ |  	| j j ¡ d¡ | j  d¡ d S )Nú!BBHé   é   é"   ú1.2.3.4ó   fooBARó    r   éZ   ó   hello, worldó   hi thereúfake reason)r8   ÚdataReceivedÚstructÚpackÚsocketÚ	inet_atonr$   ÚvalueÚclearÚassertEqualÚassertFalser   ÚassertIsNotNoner&   ÚconnectionLost)r   Úsentr   r   r   Útest_simplei   s*   ÿþýÿÿzConnectTests.test_simplec              	   C   sè   t  dddd¡t d¡ d d }t|ƒD ]}| j |¡ q| jj ¡ }| jj 	¡  |  
|t  ddd	d¡t d
¡ ¡ |  | jjj¡ |  | jj¡ | j d¡ |  
| jjj ¡ d¡ | jj d¡ |  
| jj ¡ d¡ | j d¡ dS )áu  
        If the destination IP address has zeros for the first three octets and
        non-zero for the fourth octet, the client is attempting a v4a
        connection.  A hostname is specified after the user ID string and the
        server connects to the address that hostname resolves to.

        @see: U{http://en.wikipedia.org/wiki/SOCKS#SOCKS_4a_protocol}
        r?   r@   rA   rB   ú0.0.0.1ó   fooBAZ ó
   localhost r   rF   r6   rG   rH   rI   N)rK   rL   rM   rN   r   r8   rJ   r$   rO   rP   rQ   rR   r   rS   r&   rT   ©r   ÚclientRequestÚbyterU   r   r   r   Ú test_socks4aSuccessfulResolution‚   s,   ÿþýÿ
ÿz-ConnectTests.test_socks4aSuccessfulResolutionc              	   C   óŒ   t  dddd¡t d¡ d d }t|ƒD ]}| j |¡ q| jj ¡ }|  	|t  ddd	d¡t d
¡ ¡ |  
| jjj¡ |  | jj¡ dS )ú‹
        Failed hostname resolution on a SOCKSv4a packet results in a 91 error
        response and the connection getting closed.
        r?   r@   rA   rB   rX   rY   ó   failinghost r   é[   ú0.0.0.0N©rK   rL   rM   rN   r   r8   rJ   r$   rO   rQ   r<   r   ÚassertIsNoner&   r[   r   r   r   Útest_socks4aFailedResolution°   ó    ÿþýÿ
ÿz)ConnectTests.test_socks4aFailedResolutionc              	   C   ó‚   dd„ | j _| j  t dddd¡t d¡ d d	 ¡ |  | j j 	¡ t dd
dd
¡t d¡ ¡ |  
| j jj¡ |  | j j¡ d S )Nc                 S   ó   dS ©Nr   r   ©ÚcodeÚserverr(   Úuserr   r   r   Ú<lambda>Î   ó    z0ConnectTests.test_accessDenied.<locals>.<lambda>r?   r@   rA   é’  ú10.2.3.4rD   rE   r   rb   rc   )r8   Ú	authorizerJ   rK   rL   rM   rN   rQ   r$   rO   r<   r   re   r&   r   r   r   r   Útest_accessDeniedÍ   ó    ÿþýÿ
þzConnectTests.test_accessDeniedc                 C   sx   | j  t dddd¡t d¡ d d ¡ | j j ¡  | j  d¡ |  | j j	j 
¡ d¡ | j j	j ¡  | j j	 d	¡ d S ©
Nr?   r@   rA   rB   rC   rD   rE   rG   rI   )r8   rJ   rK   rL   rM   rN   r$   rP   rQ   r&   rO   r   rT   r   r   r   r   Útest_eofRemoteÜ   s   ÿþýÿzConnectTests.test_eofRemotec                 C   sh   | j  t dddd¡t d¡ d d ¡ | j j ¡  | j  d¡ |  | j j	j 
¡ d¡ | j  d	¡ d S rv   )r8   rJ   rK   rL   rM   rN   r$   rP   rQ   r&   rO   rT   r   r   r   r   Útest_eofLocalí   s   ÿþýÿzConnectTests.test_eofLocalN)r   r   r   r"   r:   r>   rV   r^   rf   rt   rw   rx   r   r   r   r   r3   V   s    .r3   c                   @   r2   )Ú	BindTestszS
    Tests for SOCKS and SOCKSv4a bind requests using the L{SOCKSv4} protocol.
    c                 C   r4   r5   r7   r   r   r   r   r:     r;   zBindTests.setUpc              	   C   s@  | j  t dddd¡t d¡ d d ¡ | j j ¡ }| j j ¡  |  	|t ddd	d
¡t d¡ ¡ |  
| j jj¡ |  | j j¡ | j j d¡}|  |¡ tƒ |_| ¡  | j j ¡ }| j j ¡  |  	|t ddd	d¡t d¡ ¡ |  
| j jj¡ | j  d¡ |  	|j ¡ d¡ | d¡ |  	| j j ¡ d¡ | j  d¡ d S )Nr?   r@   é   rB   rC   rD   rE   r   rF   r-   r.   ©rC   éá  rc   rG   rH   rI   )r8   rJ   rK   rL   rM   rN   r$   rO   rP   rQ   rR   r   rS   r/   ÚbuildProtocolr	   r%   rT   ©r   rU   Úincomingr   r   r   rV     s>   ÿþýÿÿ
ÿ
zBindTests.test_simplec              	   C   sT  t  dddd¡t d¡ d d }t|ƒD ]}| j |¡ q| jj ¡ }| jj 	¡  |  
|t  ddd	d
¡t d¡ ¡ |  | jjj¡ |  | jj¡ | jj d¡}|  |¡ tƒ |_| ¡  | jj ¡ }| jj 	¡  |  
|t  ddd	d¡t d¡ ¡ |  | jjjd¡ | j d¡ |  
|j ¡ d¡ | d¡ |  
| jj ¡ d¡ | j d¡ dS )rW   r?   r@   rz   rB   rX   rY   rZ   r   rF   r-   r.   )r6   r|   rc   NrH   rI   )rK   rL   rM   rN   r   r8   rJ   r$   rO   rP   rQ   rR   r   rS   r/   r}   r	   r%   ÚassertIsNotrT   )r   r\   r]   rU   r   r   r   r   Útest_socks4a7  s@   ÿþýÿ
ÿ
ÿ
zBindTests.test_socks4ac              	   C   r_   )r`   r?   r@   rz   rB   rX   rY   ra   r   rb   rc   Nrd   r[   r   r   r   rf   r  rg   z&BindTests.test_socks4aFailedResolutionc              	   C   rh   )Nc                 S   ri   rj   r   rk   r   r   r   ro     rp   z-BindTests.test_accessDenied.<locals>.<lambda>r?   r@   rz   rq   rr   rD   rE   r   rb   rc   )r8   rs   rJ   rK   rL   rM   rN   rQ   r$   rO   r<   r   re   r/   r   r   r   r   rt     ru   zBindTests.test_accessDeniedc              	   C   sê   | j  t dddd¡t d¡ d d ¡ | j j ¡ }| j j ¡  | j j	 
d¡}|  |¡ tƒ |_| ¡  | j j ¡ }| j j ¡  |  |t dd	d
d	¡t d¡ ¡ |  | j jj¡ | j  d¡ |  |j ¡ d¡ |j ¡  | d¡ d S ©Nr?   r@   rz   rB   rC   rD   rE   r{   r   rF   rc   rG   rI   )r8   rJ   rK   rL   rM   rN   r$   rO   rP   r/   r}   rS   r	   r%   rQ   rR   r   r   rT   r~   r   r   r   rw   ž  s2   ÿþýÿ
ÿ
zBindTests.test_eofRemotec              	   C   sâ   | j  t dddd¡t d¡ d d ¡ | j j ¡ }| j j ¡  | j j	 
d¡}|  |¡ tƒ |_| ¡  | j j ¡ }| j j ¡  |  |t dd	d
d	¡t d¡ ¡ |  | j jj¡ | j  d¡ |  |j ¡ d¡ | j  d¡ d S r‚   )r8   rJ   rK   rL   rM   rN   r$   rO   rP   r/   r}   rS   r	   r%   rQ   rR   r   rT   r~   r   r   r   rx   ¾  s0   ÿþýÿ
ÿzBindTests.test_eofLocalc              	   C   s¨   | j  t dddd¡t d¡ d d ¡ | j j ¡ }| j j ¡  | j j	 
d¡}|  |¡ | j j ¡ }| j j ¡  |  |t dd	d
d	¡t d¡ ¡ |  | j jj¡ d S )Nr?   r@   rz   rB   rC   rD   rE   )z1.6.6.6iš  r   rb   rc   )r8   rJ   rK   rL   rM   rN   r$   rO   rP   r/   r}   re   rQ   r<   r   r~   r   r   r   Útest_badSourceÝ  s&   ÿþýÿ
ÿzBindTests.test_badSourceN)r   r   r   r"   r:   rV   r   rf   rt   rw   rx   rƒ   r   r   r   r   ry   þ   s    '; ry   )r"   rM   rK   Útwisted.internetr   r   Útwisted.internet.errorr   Útwisted.protocolsr   Útwisted.python.compatr   Útwisted.testr   Útwisted.trialr   ÚStringTransportr	   r   ÚSOCKSv4r#   ÚTestCaser3   ry   r   r   r   r   Ú<module>   s    )