2
\$\begingroup\$

This is wrapper for POP3, POP3_SSL, IMAP4, IMAP4_SSL, SMTP and SMTP_SSL classes which let them work through proxies. For creating socket with proxy I've used PySocks.

Code:

from imaplib import IMAP4, IMAP4_SSL, IMAP4_PORT, IMAP4_SSL_PORT
from poplib import POP3, POP3_SSL, POP3_PORT, POP3_SSL_PORT
from smtplib import SMTP, SMTP_SSL, SMTP_PORT, SMTP_SSL_PORT
from socks import create_connection, PROXY_TYPE_HTTP, PROXY_TYPE_SOCKS4, PROXY_TYPE_SOCKS5
import socket
import sys
from urllib.parse import urlparse
from collections import namedtuple

def imap_before_connection(self, *args):
    sys.audit("imaplib.open", self, self.host, self.port)

def smtp_before_connection(self, host, port):
    if self.debuglevel > 0:
        self._print_debug('connect: to', (host, port), self.source_address)

def pop_before_connection(*args):
    pass

DEFAULT_SETTINGS = {
    POP3_SSL: (POP3_SSL_PORT, pop_before_connection),
    IMAP4_SSL: (IMAP4_SSL_PORT, imap_before_connection),
    SMTP_SSL: (SMTP_SSL_PORT, smtp_before_connection),
    POP3: (POP3_PORT, pop_before_connection),
    IMAP4: (IMAP4_PORT, imap_before_connection),
    SMTP: (SMTP_PORT, smtp_before_connection)
}

def proxify(base_type):
    for type_, (default_port, before_connection) in DEFAULT_SETTINGS.items():
        if issubclass(base_type, type_):
            break
    else:
        raise TypeError(f"Can't proxify {base_type}")

    class Proxified(base_type):
        Proxy = namedtuple("Proxy", ("host", "port", "username", "password", "proxy_type", "rdns"))
        on_before_connection = before_connection

        @staticmethod
        def parse_proxy_string(proxy_string):
            if not proxy_string:
                return None

            parsed = urlparse(proxy_string)
            _scheme = parsed.scheme.lower()
            if _scheme in {"http", "https"}:
                proxy_type, remote_dns = PROXY_TYPE_HTTP, True
            elif _scheme in {"socks4", "socks4a"}:
                proxy_type, remote_dns = PROXY_TYPE_SOCKS4, _scheme.endswith("a")
            elif _scheme in {"socks5", "socks5h"}:
                proxy_type, remote_dns = PROXY_TYPE_SOCKS5, _scheme.endswith("h")
            else:
                raise ValueError(f'"{_scheme}" is not supported proxy type')

            return Proxified.Proxy(parsed.hostname, parsed.port, parsed.username, parsed.password,
                                   proxy_type, remote_dns)

        def __init__(self, host="", port=default_port, *args, **kwargs):
            self.proxy = self.parse_proxy_string(kwargs.pop("proxy", ""))

            super().__init__(host, port, *args, **kwargs)

        def _create_socket(self, timeout):  # used in POP3 and IMAP
            return self._get_socket(self.host, self.port, timeout)

        def _get_socket(self, host, port, timeout):  # used in SMTP
            if timeout is not None and not timeout:
                raise ValueError('Non-blocking socket (timeout=0) is not supported')

            self.on_before_connection(self, host, port)

            if not self.proxy:
                sock = socket.create_connection(
                    (host, port),
                    timeout,
                    getattr(self, "source_address", None)
                )
            else:
                sock = create_connection(  # socks.create_connection
                    (host, port),
                    timeout,
                    getattr(self, "source_address", None),
                    self.proxy.proxy_type,
                    self.proxy.host,
                    self.proxy.port,
                    self.proxy.rdns,
                    self.proxy.username,
                    self.proxy.password
                )

            ssl_context = getattr(self, "context", None) or getattr(self, "ssl_context", None)
            if ssl_context:
                return ssl_context.wrap_socket(sock, server_hostname=host)
            else:
                return sock

    return Proxified

Usage:

from imaplib import IMAP4, IMAP4_SSL
from poplib import POP3, POP3_SSL
from smtplib import SMTP, SMTP_SSL

mail, mail_pass = "email@example.com", "password"
http_proxy = "http://username:password@127.0.0.1:8080"
pop_server = "pop.example.com"
imap_server = "imap.example.com"
smtp_server = "smtp.example.com"

pop_handler = proxify(POP3)(pop_server, proxy=http_proxy)
pop_secure_handler = proxify(POP3_SSL)(pop_server, proxy=http_proxy)
imap_handler = proxify(IMAP4)(imap_server, proxy=http_proxy)
imap_secure_handler = proxify(IMAP4_SSL)(imap_server, proxy=http_proxy)
smtp_handler = proxify(SMTP)(smtp_server, proxy=http_proxy)
smtp_secure_handler = proxify(SMTP_SSL)(smtp_server, proxy=http_proxy)

It works and works good but I look on this code and .. I feel that solution is a bit cumbersome.

Could you please take a look on it and point me what I've done in "wrong" way and suggest the correct way. Thanks.

UPD.

There's also another option which looks much better (imho), but it uses mock.patch() and I'm not quite sure that it's okay to use it in production.

Code:

from poplib import POP3
from imaplib import IMAP4
from smtplib import SMTP
from functools import wraps, partial
from collections import namedtuple
from unittest.mock import patch
from urllib.parse import urlparse
from socks import create_connection, PROXY_TYPE_HTTP, PROXY_TYPE_SOCKS4, PROXY_TYPE_SOCKS5

def mock_if_proxy(method):
    @wraps(method)
    def wrapper(self, *args, **kwargs):
        if self.proxy:
            with patch("socket.create_connection",
                       partial(create_connection, proxy_type=self.proxy.proxy_type, proxy_addr=self.proxy.host,
                               proxy_port=self.proxy.port, proxy_rdns=self.proxy.rdns,
                               proxy_username=self.proxy.username, proxy_password=self.proxy.password)):
                return method(self, *args, **kwargs)
        else:
            return method(self, *args, **kwargs)
    return wrapper

def proxify(base_type):
    if not issubclass(base_type, (POP3, IMAP4, SMTP)):
        raise TypeError(f"Can't proxify {base_type}")

    class Proxified(base_type):
        # skipped Proxy and parse_proxy_string() declaration, it's the same

        if hasattr(base_type, "_create_socket"):
            _create_socket = mock_if_proxy(base_type._create_socket)
        if hasattr(base_type, "_get_socket"):
            _get_socket = mock_if_proxy(base_type._get_socket)

        def __init__(self, *args, **kwargs):
            self.proxy = self.parse_proxy_string(kwargs.pop("proxy", ""))

            super().__init__(*args, **kwargs)

    return Proxified

Maybe you could advice me a better way without mock?

\$\endgroup\$
1
\$\begingroup\$

Iteration

for type_, (default_port, before_connection) in DEFAULT_SETTINGS.items():

doesn't necessarily need a for / break: an alternative is

try:
    default_port, before_connection = next(
        vals for type_, vals in DEFAULT_SETTINGS.items()
        if issubclass(base_type, type_)
    )
except StopIteration:
    raise TypeError(f"Can't proxify {base_type}")

namedtuple

This:

  Proxy = namedtuple("Proxy", ("host", "port", "username", "password", "proxy_type", "rdns"))

isn't bad, but consider replacing it with a @dataclass having typed members. This will improve your static analysis (as much as that's meaningful in Python).

Explicit None checks

        if timeout is not None and not timeout:

is a little muddy. You have an explicit None check, but then use truthy syntax to imply a zero-check. You're better off just spelling out

if timeout == 0:

which, given this context, has the same effect.

Type hints

Particularly for signatures like

    def _get_socket(self, host: str, port: int, timeout: int) -> socket:  # used in SMTP

type hints are often even more useful than documentation in understanding what a method is supposed to do.

General

I'd recommend your first approach over using mock, which I usually only expect to see in tests.

\$\endgroup\$
5
  • \$\begingroup\$ 1. I need items() in that case, because after I break loop values of default_port and before_connection is remaining from last iteration and I use this values in class declaration. 2. I can't replace this loop with all(), I possibly can replace it with any() but it will require to iterate over loop again to search for settings. 3. I copied this explicit None check from code of imaplib and left like it is. \$\endgroup\$ Dec 11 '20 at 6:50
  • \$\begingroup\$ Doublechecked about None and found it in sources of all modules I'm working with: poplib, imaplib, smtplib. About static typing - I haven't added it because I won't use this methods at all, I'm overriding it and somewhere deep it uses by module itself, so I see no reason. \$\endgroup\$ Dec 11 '20 at 7:54
  • \$\begingroup\$ About items(), I misinterpreted what's going on - edited. \$\endgroup\$
    – Reinderien
    Dec 11 '20 at 14:49
  • \$\begingroup\$ About None - just because someone else does something bad, that doesn't mean that you have to follow suit. \$\endgroup\$
    – Reinderien
    Dec 11 '20 at 14:50
  • \$\begingroup\$ And about typing - types still have value even if they're not directly, externally visible. Your unit tests that you should be writing will benefit from them, and any sort of mypy verification will benefit from them too. \$\endgroup\$
    – Reinderien
    Dec 11 '20 at 14:51

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.