romis2012 / python-socks / 89

pending completion

Build # 89

Build Type

push

travis-ci-com

Committed by

romis2012

Commit Message

Optimize socks5 bound address reading

Run Details

6 of 6 new or added lines in 3 files covered. (100.0%)

1586 of 1597 relevant lines covered (99.31%)

3.97 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

94.12

/python_socks/_proto/socks5_async.py

import socket

from .socks5 import (
    AuthMethod,
    AddressType,
    AuthMethodsRequest,
    AuthMethodsResponse,
    AuthRequest,
    AuthResponse,
    ConnectRequest,
    ConnectResponse,
)
from .. import _abc as abc
from .._errors import ProxyError


class Socks5Proto:
    def __init__(
        self,
        stream: abc.AsyncSocketStream,
        resolver: abc.AsyncResolver,
        dest_host,
        dest_port,
        username=None,
        password=None,
        rdns=None,
    ):
        if rdns is None:
            rdns = True

        self._dest_port = dest_port
        self._dest_host = dest_host

        self._username = username
        self._password = password
        self._rdns = rdns

        self._stream = stream
        self._resolver = resolver

    async def negotiate(self):
        await self._socks_auth()
        await self._socks_connect()

    async def _socks_auth(self):
        auth_method = await self._choose_auth_method()

        # authenticate
        if auth_method == AuthMethod.USERNAME_PASSWORD:
            req = AuthRequest(username=self._username, password=self._password)

            await self._stream.write_all(bytes(req))

            res = AuthResponse(await self._stream.read_exact(2))
            res.validate()

    async def _choose_auth_method(self) -> AuthMethod:
        req = AuthMethodsRequest(
            username=self._username,
            password=self._password,
        )

        await self._stream.write_all(bytes(req))

        res = AuthMethodsResponse(await self._stream.read_exact(2))
        res.validate(request=req)
        return res.auth_method

    async def _socks_connect(self):
        req = ConnectRequest(
            host=self._dest_host,
            port=self._dest_port,
            rdns=self._rdns,
        )

        if req.need_resolve:
            _, addr = await self._resolver.resolve(req.host, family=socket.AF_UNSPEC)
            req.set_resolved_host(addr)

        await self._stream.write_all(bytes(req))

        res = ConnectResponse(await self._stream.read_exact(3))
        res.validate()

        # read remaining data (server bound address)
        # await self._stream.read()
        await self._read_bound_address()

    async def _read_bound_address(self):
        addr_type, *_ = await self._stream.read_exact(1)
        if addr_type == AddressType.IPV4:
            host = await self._stream.read_exact(4)
            host = socket.inet_ntop(socket.AF_INET, host)
        elif addr_type == AddressType.IPV6:
            host = await self._stream.read_exact(16)
            host = socket.inet_ntop(socket.AF_INET6, host)
        elif addr_type == AddressType.DOMAIN:  # pragma: no cover
            host_len, *_ = await self._stream.read_exact(1)
            host = await self._stream.read_exact(host_len)
            host = host.decode()
        else:  # pragma: no cover
            raise ProxyError('Invalid address type: {:#02X}'.format(addr_type))

        port = await self._stream.read_exact(2)
        port = int.from_bytes(port, 'big')

        return host, port

1	import socket	4✔
2
3	from .socks5 import (	4✔
4	AuthMethod,
5	AddressType,
6	AuthMethodsRequest,
7	AuthMethodsResponse,
8	AuthRequest,
9	AuthResponse,
10	ConnectRequest,
11	ConnectResponse,
12	)
13	from .. import _abc as abc	4✔
14	from .._errors import ProxyError	4✔
15
16
17	class Socks5Proto:	4✔
18	def __init__(	4✔
19	self,
20	stream: abc.AsyncSocketStream,
21	resolver: abc.AsyncResolver,
22	dest_host,
23	dest_port,
24	username=None,
25	password=None,
26	rdns=None,
27	):
28	if rdns is None:	4✔
29	rdns = True	4✔
30
31	self._dest_port = dest_port	4✔
32	self._dest_host = dest_host	4✔
33
34	self._username = username	4✔
35	self._password = password	4✔
36	self._rdns = rdns	4✔
37
38	self._stream = stream	4✔
39	self._resolver = resolver	4✔
40
41	async def negotiate(self):	4✔
42	await self._socks_auth()	4✔
43	await self._socks_connect()	4✔
44
45	async def _socks_auth(self):	4✔
46	auth_method = await self._choose_auth_method()	4✔
47
48	# authenticate
49	if auth_method == AuthMethod.USERNAME_PASSWORD:	4✔
50	req = AuthRequest(username=self._username, password=self._password)	4✔
51
52	await self._stream.write_all(bytes(req))	4✔
53
54	res = AuthResponse(await self._stream.read_exact(2))	4✔
55	res.validate()	4✔
56
57	async def _choose_auth_method(self) -> AuthMethod:	4✔
58	req = AuthMethodsRequest(	4✔
59	username=self._username,
60	password=self._password,
61	)
62
63	await self._stream.write_all(bytes(req))	4✔
64
65	res = AuthMethodsResponse(await self._stream.read_exact(2))	4✔
66	res.validate(request=req)	4✔
67	return res.auth_method	4✔
68
69	async def _socks_connect(self):	4✔
70	req = ConnectRequest(	4✔
71	host=self._dest_host,
72	port=self._dest_port,
73	rdns=self._rdns,
74	)
75
76	if req.need_resolve:	4✔
77	_, addr = await self._resolver.resolve(req.host, family=socket.AF_UNSPEC)	4✔
78	req.set_resolved_host(addr)	4✔
79
80	await self._stream.write_all(bytes(req))	4✔
81
82	res = ConnectResponse(await self._stream.read_exact(3))	4✔
83	res.validate()	4✔
84
85	# read remaining data (server bound address)
86	# await self._stream.read()
87	await self._read_bound_address()	4✔
88
89	async def _read_bound_address(self):	4✔
90	addr_type, *_ = await self._stream.read_exact(1)	4✔
91	if addr_type == AddressType.IPV4:	4✔
92	host = await self._stream.read_exact(4)	4✔
93	host = socket.inet_ntop(socket.AF_INET, host)	4✔
94	elif addr_type == AddressType.IPV6:	×
95	host = await self._stream.read_exact(16)	×
96	host = socket.inet_ntop(socket.AF_INET6, host)	×
97	elif addr_type == AddressType.DOMAIN: # pragma: no cover
98	host_len, *_ = await self._stream.read_exact(1)
99	host = await self._stream.read_exact(host_len)
100	host = host.decode()
101	else: # pragma: no cover
102	raise ProxyError('Invalid address type: {:#02X}'.format(addr_type))
103
104	port = await self._stream.read_exact(2)	4✔
105	port = int.from_bytes(port, 'big')	4✔
106
107	return host, port	4✔

romis2012 / python-socks / 89

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous