1import curio
2import curio.io
3
4from ..._errors import ProxyConnectionError, ProxyTimeoutError
5from ..._proto.http_async import HttpProto
6from ..._proto.socks4_async import Socks4Proto
7from ..._proto.socks5_async import Socks5Proto
8from ._stream import CurioSocketStream
9from ._resolver import Resolver
10from ._connect import connect_tcp
11
12from ... import _abc as abc
13
14DEFAULT_TIMEOUT = 60
15
16
17class CurioProxy(abc.AsyncProxy):
18    def __init__(
19        self,
20        proxy_host: str,
21        proxy_port: int,
22    ):
23        self._proxy_host = proxy_host
24        self._proxy_port = proxy_port
25
26        self._dest_host = None
27        self._dest_port = None
28        self._timeout = None
29
30        self._stream = None
31        self._resolver = Resolver()
32
33    async def connect(
34        self,
35        dest_host: str,
36        dest_port: int,
37        timeout: float = None,
38        _socket=None,
39    ) -> curio.io.Socket:
40        if timeout is None:
41            timeout = DEFAULT_TIMEOUT
42
43        self._dest_host = dest_host
44        self._dest_port = dest_port
45        self._timeout = timeout
46
47        try:
48            return await curio.timeout_after(self._timeout, self._connect, _socket)
49        except OSError as e:
50            await self._close()
51            msg = 'Could not connect to proxy {}:{} [{}]'.format(
52                self._proxy_host,
53                self._proxy_port,
54                e.strerror,
55            )
56            raise ProxyConnectionError(e.errno, msg) from e
57        except curio.TaskTimeout as e:
58            await self._close()
59            raise ProxyTimeoutError('Proxy connection timed out: %s' % self._timeout) from e
60        except Exception:
61            await self._close()
62            raise
63
64    async def _connect(self, _socket=None):
65        if _socket is None:
66            _socket = await connect_tcp(
67                host=self._proxy_host,
68                port=self._proxy_port,
69            )
70        self._stream = CurioSocketStream(_socket)
71        await self._negotiate()
72        return _socket
73
74    async def _negotiate(self):
75        raise NotImplementedError()
76
77    async def _close(self):
78        if self._stream is not None:
79            await self._stream.close()
80
81    @property
82    def proxy_host(self):
83        return self._proxy_host
84
85    @property
86    def proxy_port(self):
87        return self._proxy_port
88
89
90class Socks5Proxy(CurioProxy):
91    def __init__(self, proxy_host, proxy_port, username=None, password=None, rdns=None):
92        super().__init__(proxy_host=proxy_host, proxy_port=proxy_port)
93        self._username = username
94        self._password = password
95        self._rdns = rdns
96
97    async def _negotiate(self):
98        proto = Socks5Proto(
99            stream=self._stream,
100            resolver=self._resolver,
101            dest_host=self._dest_host,
102            dest_port=self._dest_port,
103            username=self._username,
104            password=self._password,
105            rdns=self._rdns,
106        )
107        await proto.negotiate()
108
109
110class Socks4Proxy(CurioProxy):
111    def __init__(self, proxy_host, proxy_port, user_id=None, rdns=None):
112        super().__init__(proxy_host=proxy_host, proxy_port=proxy_port)
113        self._user_id = user_id
114        self._rdns = rdns
115
116    async def _negotiate(self):
117        proto = Socks4Proto(
118            stream=self._stream,
119            resolver=self._resolver,
120            dest_host=self._dest_host,
121            dest_port=self._dest_port,
122            user_id=self._user_id,
123            rdns=self._rdns,
124        )
125        await proto.negotiate()
126
127
128class HttpProxy(CurioProxy):
129    def __init__(self, proxy_host, proxy_port, username=None, password=None):
130        super().__init__(proxy_host=proxy_host, proxy_port=proxy_port)
131        self._username = username
132        self._password = password
133
134    async def _negotiate(self):
135        proto = HttpProto(
136            stream=self._stream,
137            dest_host=self._dest_host,
138            dest_port=self._dest_port,
139            username=self._username,
140            password=self._password,
141        )
142        await proto.negotiate()
143