Добавляем кастомные HTTP заголовки в Selenium-тесты на Python за 4 простых шага.

Статистика показывает, что работа с кастомными заголовками в Selenium-тестах является весьма распространенной проблемой. В одной из моих предыдущих статей, я показывал один из способов решения такой задачи при помощи библиотеки Browsermob-Proxy в тестах на языке Java. Однако, люди, привыкшие программировать на языке Python, также часто встречают необходимость работы с заголовками запросов, отправляемых браузером, находящимся под контролем Selenium.

Для таких людей существует как минимум два способа решения упомянутой проблемы. Первый способ - использовать клиентскую обертку для BrowserMob-Proxy. Недостаток такого подхода заключается в том, что вам всё же придется иметь дело с JRE. Прокси будет постоянно запущен, а использовать вы его будете, обращаясь к его REST-API, через клиентскую библиотеку на Python. Другим способом является использование чисто питоновского прокси-сервера. Об этому мы и будем разговаривать в сегодняшней статье.

Описание примера: Поддержка базовой авторизации в Selenium-тесте.

Чтобы быть более конкретным, я покажу пример изменения заголовка запросов для решения задачи доступа к странице, скрытой механизмом базовой аутентификации. Тестовую страницу, с включенным подобным механизмом, вы можете найти тут. Для ручной проверки используйте webelement в качестве имени пользователя и click в качестве пароля.

Я разобью решение на 4 несложных шага: добавление необходимой библиотеки в ваше окружение, реализация кастомной части, написание теста и добавление поддержки протокола HTTPs.

Добавление поддержки HTTPs потребует использования инструмента openssl (некоторые дистрибутивы Linux, такие как Ubuntu, имеют данную утилиту в своих репозиториях). Вы воспользуетесь openssl один раз для включения механизма MITM, который позволит расшифровывать и модифицировать TLS-трафик.

Шаг 1 - Установка Proxy.Py

Proxy.Py - прокси сервер, написанный на чистом Питоне, поддерживающий множество удобных функций. Гораздо больший набор, нежели необходимый нам для нашей задачи. Этот сервер может запускаться как сам по себе, так и в составе стороннего кода (последняя возможность для нас особенно ценна, т.к. мы будем запускать и отключать его из наших автоматизированных тестов на Selenium).

Возможности сервера могут быть кастомизированы. Библиотека поддерживает расширяемость через концепцию плагинов. Поэтому, нашим следующим шагом станет реализация такого плагина. А пока, нам необходимо установить Proxy.Py в наше окружение:

pip install --upgrade proxy.py

Теперь вы можете использовать библиотеку в вашем коде на Python.

Иногда в системе может быть установлено несколько версий Python. Убедитесь, что вы используете правильную версию pip, соответствующую версии Питона, которую планируете в дальнейшем использовать. Например, в моей системе установлен Python версий 2.7 и 3.8. Так, команда pip устанавливает пакеты для версии 2.7, а команда pip3 - для версии 3.8.

Шаг 2 - Готовим кастомный плагин

После установки Proxy.Py мы можем приступить к нашей кастомизиции. Нам необходимо создать плагин для прокси, который бы модифицировал HTTP-заголовки, проходящие через прокси-сервер. Наша цель - добавить специальный заголовок, хранящий параметры нашей учетной записи (более подробно о "протокольной" стороне вопроса вы можете узнать либо перейдя по ссылке, которую я вынес во вступление, либо тут).

Итак, создайте файл, где мы разместим код нашего плагина. Назовем этот файл header_modifier.py. Поместите туда следующий код:

from proxy.http.proxy import HttpProxyBasePlugin
from proxy.http.parser import HttpParser
from typing import Optional
import base64


class BasicAuthorizationPlugin(HttpProxyBasePlugin):
    """Modifies request headers."""

    def before_upstream_connection(
            self, request: HttpParser) -> Optional[HttpParser]:
        return request

    def handle_client_request(
            self, request: HttpParser) -> Optional[HttpParser]:
        basic_auth_header = 'Basic ' + base64.b64encode('webelement:click'.encode('utf-8')).decode('utf-8')
        request.add_header('Authorization'.encode('utf-8'), basic_auth_header.encode('utf-8'))
        return request

    def on_upstream_connection_close(self) -> None:
        pass

    def handle_upstream_chunk(self, chunk: memoryview) -> memoryview:
        return chunk

Здесь мы создали наш кастомный плагин, расширяющий уже имеющийся плагин HttpProxyBasePlugin. Мы переопределили метод handle_client_request в котором мы добавляем в реквест аутентификационный заголовок, в соответствие со спецификацией RFC7235.

Шаг 3 - Добавляем поддержку HTTPs

Кода, которым мы располагаем к данному шагу, было бы достаточно, если мы хотели поддерживать HTTPs в наших тестах. Дело в том, что когда вы используете HTTPs, ваш бразуер шифрует исходящие сообщения, таким образом, изменить в них что-либо не возможно. Решением этой проблемы станет включение MITM-модуля, который встраивается во взаимодействие браузера и приложения. Этот модуль сам становится "клиентом" веб-приложения, расшифровывает трафик от сервера, как это бы сделал браузер, и перезашифровывает его при помощи собственного сертификата.

Proxy.Py умеет выпускать собственные сертификаты для сайтов "на лету". Для того, чтобы это стало возможным, нам необходимо подготовить три файла:

  • Приватный ключ, который будет использоваться для подписания выпущенных серверных сертификатов

  • Сертификат, который будет выступать в качестве корневого

  • Приватный ключ, который будет использоваться для генерации серверных сертификатов.

Здесь и далее я предполагаю, что вы создадите отдельный каталог, в котором будут храниться файлы ключей. Допустим таким каталогом станет /test/mitm. Находясь в каталоге, сделайте следующее:

1) Сгенерируйте приватный ключ для подписи сертификатов:

openssl genrsa -out wec-ca.key 2048

2) Сгенерируйте корневой сертификат:

openssl req -x509 -new -nodes -key wec-ca.key -sha256 -days 1825 -out wec-ca.pem

3) Сгенерируйте приватный ключ для генерации серверных сертификатов

openssl genrsa -out wec-signing.key 2048

К этому моменту в нашем каталоге должно находиться три файла. Если это так, мы можем переходить к финальному шагу - написанию теста.

Шаг 4 - Пишем тест

На этом финальном шаге мы объединим все компоненты нашего решения. Мы напишем Selenium-тест, который будет запускать прокси-сервер и открывать HTTPs-страницу, скрытую за базовой аутентификацией, используя запущенный прокси-сервер.

Я подразуменваю, что у вас установлены клиентские библиотеки Selenium на языке Python, а также установлен сам webDriver. В пример я не использую какие-либо тестовые фреймворки, так что наличие Proxy.Py and Selenium, а также стандартный питоновских библиотек должно хватить.

Итак, создайте файл main.py рядом с уже созданным файлом плагина. Добавьте следующий код в новый файл:

from selenium import webdriver
import proxy
import header_modifier
selenium_proxy = webdriver.Proxy()


def run_test():
    from selenium.webdriver import DesiredCapabilities
    capabilities = DesiredCapabilities.FIREFOX
    selenium_proxy.add_to_capabilities(capabilities)
    driver = webdriver.Firefox(capabilities=capabilities)
    driver.get('https://www.webelement.click/stand/basic?lang=en')
    assert driver.find_element_by_css_selector('.post-body h2').text == 'You have authorized successfully!'
    driver.quit()


if __name__ == '__main__':
    from proxy.common import utils
    proxy_port = utils.get_available_port()
    with proxy.start(
            ['--host', '127.0.0.1',
             '--port', str(proxy_port),
             '--ca-cert-file', '/test/mitm/wec-ca.pem',
             '--ca-key-file', '/test/mitm/wec-ca.key',
             '--ca-signing-key-file', '/test/mitm/wec-signing.key'],
            plugins=
            [b'header_modifier.BasicAuthorizationPlugin',
             header_modifier.BasicAuthorizationPlugin]):
        from selenium.webdriver.common.proxy import ProxyType
        selenium_proxy.proxyType = ProxyType.MANUAL
        selenium_proxy.httpProxy = '127.0.0.1:' + str(proxy_port)
        selenium_proxy.sslProxy = '127.0.0.1:' + str(proxy_port)
        print('Proxy address: ' + selenium_proxy.httpProxy)
        run_test()

Давайте подробнее взглянем на этот код. Ниже представлен метод, запускающий тест. Он запускается в самом конце, когда вся необходимые настроечные процедуры уже проделаны. Тест связывает браузер Firefox с нашим прокси и обращается к странице, скрытой механизмом базовой аутентификации. Затем тест ищет некоторый текст на странице. Тест заканчивается успешно только в том случае, если аутентификация пройдена.

def run_test():
    from selenium.webdriver import DesiredCapabilities
    capabilities = DesiredCapabilities.FIREFOX
    selenium_proxy.add_to_capabilities(capabilities)
    driver = webdriver.Firefox(capabilities=capabilities)
    driver.get('https://www.webelement.click/stand/basic?lang=en')
    assert driver.find_element_by_css_selector('.post-body h2').text == 'You have authorized successfully!'
    driver.quit()

Затем мы переходим к основной части:

from proxy.common import utils
proxy_port = utils.get_available_port()
with proxy.start(
        ['--host', '127.0.0.1',
         '--port', str(proxy_port),
         '--ca-cert-file', '/test/mitm/wec-ca.pem',
         '--ca-key-file', '/test/mitm/wec-ca.key',
         '--ca-signing-key-file', '/test/mitm/wec-signing.key'],
        plugins=
        [b'header_modifier.BasicAuthorizationPlugin',
         header_modifier.BasicAuthorizationPlugin]):

Эта часть запускается в первую очередь. Начинаем мы с того, что получаем значение свободного порта, на котором наш прокси будет работать. Запускаем прокси-сервер с параметрами, задающими настройку текущего запуска. Некоторые из этих параметров, активируют MITM-модуль. Значения этих параметров указывают на сгенеренные нами при помощи openssl файлы.

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

В принципе, это всё. Если вы запустите код теста, вы не должны получить эксепшенов. Вы можете поменять ожидаемый со страницы текст, для того чтобы убедиться, что аутентификация проходит корректно, а значит, заголовок действительно добавляется к запросам.

Надеюсь, представленные примеры достаточно репрезентативны и описаны простым языком. Если у вас всё же остались вопросы, направляйте их мне через форму обратной связи. Я обязательно отвечу, а также постараюсь исправить недочеты в статье.