The way to pass proxy into client

[bahkadomos]

Some people asked how to connect to IMAP via proxy. I suggest to use python-socks package for this purpose. Thanks to the author for helping with it.

import asyncio
from dataclasses import dataclass
import ssl
from typing import Callable, Literal

from aioimaplib.aioimaplib import (
    IMAP4,
    IMAP4ClientProtocol,
)
from python_socks.async_.asyncio import Proxy
from python_socks._errors import ProxyConnectionError, ProxyError, ProxyTimeoutError
from python_socks._types import ProxyType


class EmailError(Exception):
    pass


class EmailProxyError(EmailError):
    def __init__(self) -> None:
        super().__init__(f'IMAP proxy error')


@dataclass(kw_only=True)
class ProxySchema:
    proto: Literal['http', 'socks4', 'socks5']
    host: str
    port: int
    username: str | None
    password: str | None


class ImapProxyClient(IMAP4):
    def __init__(
        self,
        *,
        host: str,
        port: int,
        loop: asyncio.AbstractEventLoop | None = None,
        timeout: float = 30,
        proxy: ProxySchema | None = None
    ):
        self._proxy = proxy
        self._loop = loop or asyncio.get_running_loop()
        super().__init__(host, port, self._loop, timeout)

    @property
    def sock_type(self) -> ProxyType | None:
        if not hasattr(self, '_sock_type'):
            if self._proxy is None:
                self._sock_type = None
            else:
                protos = {
                    'http': ProxyType.HTTP,
                    'socks4': ProxyType.SOCKS4,
                    'socks5': ProxyType.SOCKS5
                }
                self._sock_type = protos.get(self._proxy.proto.lower())
        return self._sock_type

    def create_client(
        self,
        host: str,
        port: int,
        loop: asyncio.AbstractEventLoop | None = None,
        conn_lost_cb: Callable[[Exception | None], None] = None, # type: ignore
        ssl_context: ssl.SSLContext | None = None
    ) -> None:
        self.protocol = IMAP4ClientProtocol(self._loop, conn_lost_cb)
        if self._proxy and self.sock_type:
            self._loop.create_task(self._proxy_connect(loop or self._loop, lambda: self.protocol, ssl_context))
        else:
            self._loop.create_task(self._loop.create_connection(
                lambda: self.protocol,
                host,
                port,
                ssl=ssl_context
            ))

    async def _proxy_connect(
        self,
        loop: asyncio.AbstractEventLoop,
        protocol_factory,
        ssl_context: ssl.SSLContext | None = None
    ):
        if self._proxy and self.sock_type:
            proxy = Proxy.create(
                proxy_type=self.sock_type,
                host=self._proxy.host,
                port=self._proxy.port,
                username=self._proxy.username,
                password=self._proxy.password,
                loop=loop
            )
            try:
                sock = await proxy.connect(self.host, self.port, timeout=self.timeout)
            except (ProxyError, ProxyConnectionError, ProxyTimeoutError):
                raise EmailProxyError()
            await loop.create_connection(
                protocol_factory,
                sock=sock,
                ssl=ssl_context,
                server_hostname=self.host if ssl_context else None
            )


class ImapSslProxyClient(ImapProxyClient):
    def __init__(
        self,
        *,
        host: str,
        port: int,
        timeout: float = 30,
        proxy: ProxySchema | None = None
    ):
        super().__init__(
            host=host,
            port=port,
            timeout=timeout,
            proxy=proxy
        )

    def create_client(
        self,
        host: str,
        port: int,
        loop: asyncio.AbstractEventLoop | None = None,
        conn_lost_cb: Callable[[Exception | None], None] = None, # type: ignore
        ssl_context: ssl.SSLContext | None = None
    ) -> None:
        if ssl_context is None:
            ssl_context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH)
        super().create_client(host, port, loop or self._loop, conn_lost_cb, ssl_context)


async def test():
    proxy = ProxySchema(proto='socks5', host='127.0.0.1', port=123, username=None, password=None)
    client = ImapSslProxyClient(host='host.imap.com', port=993, timeout=10, proxy=proxy)
    await client.wait_hello_from_server()
    await client.login('user', 'password')
    await client.select()
    res = await client.search('(ALL)')
    print(res)


asyncio.run(test())