旧代码拾遗：如何在 Python 代码中修改 DNS 解析

# -*- coding: utf-8 -*-
import contextlib
import importlib
import threading
from typing import Callable, Dict, Optional

from urllib3.util import connection


class CustomLocalDnsResolver(threading.local):
    """支持线程级自定义 Dns 记录
    """

    def __init__(self, dns_map: Optional[dict]=None):
        # 线程保存各自 dns_map，但是访问入口均为 dns_map
        self.dns_map = dns_map or {}


def get_patch_create_connection_with_dns(dns_resolver) -> Callable:
    """simply get patched create_connection"""

    # 保留原方法
    _orig_create_connection = getattr(importlib.import_module('urllib3.util.connection'), 'create_connection')

    def patched_create_connection(address, *args, **kwargs):
        """在 urllib3's create_connection 流程前解析 address"""
        domain, port = address
        # 当 _local_dns.dns_map 为空，对正常流程无影响
        host = dns_resolver.dns_map.get(domain, domain)
        return _orig_create_connection((host, port), *args, **kwargs)

    return patched_create_connection


_local_dns = CustomLocalDnsResolver()
# patch 全局 create_connection
connection.create_connection = get_patch_create_connection_with_dns(_local_dns)


@contextlib.contextmanager
def update_local_dns_once(dns_map: Dict):
    """一次性修改线程 dns 解析"""
    _local_dns.dns_map = dns_map
    yield
    _local_dns.dns_map = {}


# 具体的使用场景
with update_local_dns_once({"kubernetes": "192.168.1.1"}):
    # 可以在该 context 中请求外部系统