-
Notifications
You must be signed in to change notification settings - Fork 89
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Added SOCKSConnection creation into connection pool. Moved SOCKSConne…
…ction to the dedicated file to hide this class under socksio import check
- Loading branch information
Showing
9 changed files
with
360 additions
and
313 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,121 @@ | ||
| from ssl import SSLContext | ||
|
|
||
| from socksio import socks5 | ||
|
|
||
| from .._backends.auto import AsyncBackend, AsyncSocketStream | ||
| from .._exceptions import ProxyError | ||
| from .._types import Origin, Socks5ProxyCredentials, TimeoutDict | ||
| from .._utils import get_logger | ||
| from .connection import AsyncHTTPConnection | ||
|
|
||
| logger = get_logger(__name__) | ||
|
|
||
|
|
||
| class AsyncSOCKSConnection(AsyncHTTPConnection): | ||
| def __init__( | ||
| self, | ||
| origin: Origin, | ||
| http2: bool = False, | ||
| uds: str = None, | ||
| ssl_context: SSLContext = None, | ||
| socket: AsyncSocketStream = None, | ||
| local_address: str = None, | ||
| backend: AsyncBackend = None, | ||
| *, | ||
| proxy_origin: Origin, | ||
| proxy_credentials: Socks5ProxyCredentials = None, | ||
| ): | ||
| assert proxy_origin[0] in (b"socks5",) | ||
|
|
||
| super().__init__( | ||
| origin, http2, uds, ssl_context, socket, local_address, backend | ||
| ) | ||
| self.proxy_origin = proxy_origin | ||
| self.proxy_connection = socks5.SOCKS5Connection() | ||
| self.proxy_credentials = proxy_credentials | ||
|
|
||
| async def _open_socket(self, timeout: TimeoutDict = None) -> AsyncSocketStream: | ||
| _, proxy_hostname, proxy_port = self.proxy_origin | ||
| scheme, hostname, port = self.origin | ||
| ssl_context = self.ssl_context if scheme == b"https" else None | ||
| timeout = timeout or {} | ||
|
|
||
| try: | ||
| proxy_socket = await self.backend.open_tcp_stream( | ||
| proxy_hostname, | ||
| proxy_port, | ||
| None, | ||
| timeout, | ||
| local_address=self.local_address, | ||
| ) | ||
|
|
||
| await self._auth_proxy(proxy_socket, timeout) | ||
| await self._connect_through_proxy(proxy_socket, hostname, port, timeout) | ||
|
|
||
| if ssl_context: | ||
| proxy_socket = await proxy_socket.start_tls( | ||
| hostname, ssl_context, timeout | ||
| ) | ||
|
|
||
| return proxy_socket | ||
| except Exception: # noqa: PIE786 | ||
| self.connect_failed = True | ||
| raise | ||
|
|
||
| async def _auth_proxy( | ||
| self, socket: AsyncSocketStream, timeout: TimeoutDict | ||
| ) -> None: | ||
| auth_request = socks5.SOCKS5AuthMethodsRequest( | ||
| [ | ||
| socks5.SOCKS5AuthMethod.NO_AUTH_REQUIRED, | ||
| socks5.SOCKS5AuthMethod.USERNAME_PASSWORD, | ||
| ] | ||
| ) | ||
|
|
||
| self.proxy_connection.send(auth_request) | ||
|
|
||
| bytes_to_send = self.proxy_connection.data_to_send() | ||
| await socket.write(bytes_to_send, timeout) | ||
|
|
||
| data = await socket.read(1024, timeout) | ||
| auth_ev = self.proxy_connection.receive_data(data) | ||
|
|
||
| if auth_ev.method == socks5.SOCKS5AuthMethod.USERNAME_PASSWORD: # type: ignore | ||
| if self.proxy_credentials is None: | ||
| raise ProxyError( | ||
| "This proxy requires auth, but you didn't set user/password" | ||
| ) | ||
|
|
||
| user, password = self.proxy_credentials | ||
| user_password_request = socks5.SOCKS5UsernamePasswordRequest(user, password) | ||
| self.proxy_connection.send(user_password_request) | ||
| await socket.write(self.proxy_connection.data_to_send(), timeout) | ||
| user_password_response = await socket.read(2048, timeout) | ||
|
|
||
| user_password_event = self.proxy_connection.receive_data( | ||
| user_password_response | ||
| ) | ||
|
|
||
| if not user_password_event.success: # type: ignore | ||
| raise ProxyError("Invalid user/password provided to proxy auth") | ||
|
|
||
| async def _connect_through_proxy( | ||
| self, | ||
| socket: AsyncSocketStream, | ||
| hostname: bytes, | ||
| port: int, | ||
| timeout: TimeoutDict, | ||
| ) -> None: | ||
| connect_request = socks5.SOCKS5CommandRequest.from_address( | ||
| socks5.SOCKS5Command.CONNECT, (hostname, port) | ||
| ) | ||
|
|
||
| self.proxy_connection.send(connect_request) | ||
| bytes_to_send = self.proxy_connection.data_to_send() | ||
|
|
||
| await socket.write(bytes_to_send, timeout) | ||
| data = await socket.read(1024, timeout) | ||
| event = self.proxy_connection.receive_data(data) | ||
|
|
||
| # development only assert | ||
| assert event.reply_code == socks5.SOCKS5ReplyCode.SUCCEEDED # type: ignore |
Oops, something went wrong.