Пишем прокси сервер DNS запросов с кэшированием на Python.

В этой статье мы напишем простой код для проксирования сетевых запросов. То есть наша программа будет принимать запросы которые адресуються другим, отправлять их адресату, получать ответ и отправлять их получателю. Такое вот промежуточное звено или прокси. Все это дело мы реализуем с помощью встроенной библиотеке socket. Так же реализуем кэширование запросов с помощью декораторов.

Конструктивно декоратор в Python представляет собой некоторую функцию, аргументом которой является другая функция. Декоратор предназначен для добавления дополнительного функционала к данной функции без изменения содержимого последней.

Такая обертка функции, которую потом можно использовать как функцию.

Step 1

Импортируем библиотеки и объявляем глобальные переменные:

import socket

local_host = '127.0.0.1' # Адрес на котором будем принимать соединение
local_port = 53 # Порт
remote_host = '192.168.0.1' # Адрес куда будем отправлять
remote_port = 53
_cache = {}  # Словарь где будем хранить кэш
_DEBUG = True

Step 2

Функция получения данных :

# Функция принимает как аргумент объект типа socket
def receive_from(_socket): 
    _socket.settimeout(1) # TimeOut
    try:
        data, addres = _socket.recvfrom(512)
    except:
        data = ''   # Если ни каких данных не пришло, возвращаем пустые данные
        addres = ('', 0)
        return data, addres
    return data, addres

Следующая функция будет пересылать наши запросы DNS серверу:

def dns_receive_remore(local_buffer, local_addr, remote_socket):
# Передаем данные которые получили локально, адрес отправителя, и сокет для отправки данных.
# Тут надо понять что у нас есть два сокета. Один для получения данных от пользователей и отправки ему, другой для отправки DNS сервер и получения данных от него.

    if len(local_buffer) and len(local_addr[0]): # Если локальные данные не пустые
        try:
            remote_socket.sendto(local_buffer, (remote_host, remote_port)) # Отправляем данные днс серверу
        except:
            print('[!]Can not send DNS to remote.')
        remote_buffer, remore_addr = receive_from(remote_socket) # Вызываем процедуру получение данных.
        if len(remote_buffer):
            return remote_buffer
    return None

Теперь перейдем к mein функции или где все будет работать:

def server_loop(local_host, local_port):
    global _DEBUG
    server = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        server.bind((local_host, local_port)) # Запускаем локальный сокет для получения запросов
    except:
        print("[!!] Failed to listen on %s:%d" % (local_host, local_port))
        print("[!!] Check for other listening sockets or correct permissions.")
    print("[*] Listening on %s:%d" % (local_host, local_port))

# Создаем сокет для отправки запросов DNS серверу
    remote_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)

# Тут мы говорим что декоратор memoize будет использован для функции dns_receive_remore
# и вызывать мы его будет как функцию cache
    cache = memoize(dns_receive_remore)
    while True:

# Получаем локальный запрос
        local_buffer, local_addr = receive_from(server)

# Получаем ответ от DNS сервера или из кэша
        remote_buffer = cache(local_buffer, local_addr, remote_socket) 
        if remote_buffer is not None:
            server.sendto(remote_buffer, local_addr) # Отправляем ответ локальному пользователю.
            if _DEBUG:
                print('Send localhost %d bytes' % len(remote_buffer)) # маленький дебаг, что запросы получаются и отправляются
                _DEBUG = False

Декоратор

Теперь давайте разберемся с функцией декоратора:

# Функция как параметр принимает другую функцию
def memoize(func): 
    """ Декоратор для обработки кеша запроса функции."""
        name = func.__name__  #Получаем имя переданной функции
    def wrapper(*args, **kwargs):

#Первый аргумент у нас local_buffer из которого мы убираем первые два байта(ID DNS запроса)
# В DNS запросе и ответе первые два байта отведено под ID который создает клиент во время отправки запроса DNS серверу.
# Для каждого запроса и ответа он разный.
        dns_not_id_header = args[0][2:]
 
        _id = args[0][:2] #Сохраняем ID запроса что бы вернуть локальному пользователю
        key = (name, dns_not_id_header, frozenset(kwargs.items())) #создаем ключ для словаря
        if key in _cache:
            if _cache[key] is not None:
                print('[*] Received cache DNS %d bytes from localhost' % len(_cache[key]))
                print('[?] Len caches: ', len(_cache))
            return _id + _cache[key] # Если ключ есть в кэше возвращаем ID запроса + значение
        result = func(*args, **kwargs) #Если нету в кеше выполняем переданную функцию с переданными аргументами
        if result is not None:
            _cache[key] = result[2:] #Сохраняем для ключа значение в кэш
        return result
    return wrapper

Step 3

Ну а теперь запускаем нашу функцию. Eсли что то по коду не понятно, пишите комментарии, отвечу и исправлю в статье. Успехов )))

server_loop(local_host, local_port)