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