Add custom HTTP headers in Selenium Webdriver tests in Python in four simple steps

Statistics shows that working with custom headers in Selenium automated tests is one of the most popular problems that visitors of my blog try to solve. In my older post I demonstrated how to address such issue with the help of Browsermob-Proxy library in Java language. However people who code in Python also often need the tool or framework which could help to manipulate headers in requests which your browser (that is under Selenium control) sends to web application server.

There are at least two ways how one can add header manipulation support to their Selenium test. One is to use Python wrapper for BrowserMob-Proxy standalone server. The downside is that you need to run Java at your host because the proxy would still need some JRE. Another one is to use pure Python proxy and this is what we’re going to take a closer look at in this article.

Example description: Support basic authentication in your Selenium test

To be more concrete I’m going to show you header manipulation problem solving in the example where we will need to test the web page behind basic authentication mechanism. This is the page where such the authentication is used. You can check it manually. Use webelement as your user name and click as your password.

I break my solution down to 4 simple steps such as adding proxy library to your environment, implementing the custom part of proxy, writing the test and taking couple of simple steps to support HTTPs protocol in your tests.

Adding HTTPs support will require openssl tool (some Linux distributions like Ubuntu have openssl available from their software repositories). You will use it once to enable MITM mechanism that would allow to decrypt TLS traffic.

Step 1 - Install Proxy.Py

Proxy.Py is the pure Python proxy that supports lots of handy features. Much much more than we need to address our problem. It can be run in standalone mode and in embedded mode (the latter is what we actually need for our Selenium automated tests).

The proxy features can be customized in different ways. The tool uses the concept of plugins to support extensibility. Thus in the next step we’ll implement such plugin. Yet we have to install it to our environment:

pip install --upgrade proxy.py

Now you can use the library in your Python code.

Sometimes people have several version of Python installed in their systems. Make sure you’re using the right version of pip tool. For example I have both Python 2.7 and Python 3.8 installed. Thus pip installs packages for Python 2.7 and pip3 installs packages for Python 3.8.

Step 2 - Prepare custom plugin

After we have installed Proxy.Py we can proceed with our customization. We need to create plugin for our proxy. That plugin would modify the headers of HTTP requests which come through the proxy. Our goal is to add authorization header in a special manner (the details you can find in either the post I put the link at the introduction or check even more detailed description here).

So, create a file where we put the code of our plugin. Say it would be header_modifier.py. Add the following code to that file:

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

Here we create our custom plugin that extends HttpProxyBasePlugin. Our goal is to override handle_client_request method where we set up authorization header according to the corresponding specification and add that header to the processed request.

Step 3 - Add HTTPs support

The above code would basically be all we need if we wouldn’t need to test HTTPs applications. The thing is that when you use HTTPs, your browser encrypts the outgoing messages so that you cannot inject anything inside them. The solution is to enable MITM module that would pretend to be a client for an application, decrypt that traffic, and then pass it to your browser. The same is applicable for the opposite direction. Your browser encrypts the messages that your proxy now can decrypt, modify and repack in TLS for the target application.

Proxy.Py MITM module can produce fake certificates for the sites you access on the fly. To do that it becomes a sort of Certification Authority and issues the certificate for the current resource. There are three files we need to prepare to make that process working. They are:

  • Private key that would be using as CA key for signing the generated certificate

  • Certificate that would be using as CA root certificate

  • Private key that would be using to generate server certificate that then would be signed with CA private key.

Here and after I assume that you have created a dedicated folder for those files. Say that folder is /test/mitm. So when you’re inside that folder, do the following:

1) Generate CA private key file:

openssl genrsa -out wec-ca.key 2048

2) Generate CA root certificate:

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

3) Generate private key file for server certificate generation:

openssl genrsa -out wec-signing.key 2048

Now you should have three files and we can proceed to the final step - writing a test.

Step 4 - Implement your test

At this final step we combine all the components of our solution. We’ll write Selenium test that would start our proxy and access HTTPs page hidden behind basic authentication through that proxy.

I assume that you have Selenium Python bindings and WebDriver installed. I also assume that you are aware how to run your browser accepting insecure certificates. I’m not using any test framework here so except of Proxy.Py and Selenium, the standard Python package would be enough.

So create the file main.py next to already created plugin file. Add the following code there:

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()

Let’s examine this code a bit. Below is the routine that starts the test. It is started at that very end when all the required things are configured. The test sets up proxy for Firefox browser and opens a page that is behind basic authentication. Then it looks up the text on the page. The test passes only if the authentication mechanism is bypassed.

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()

Then we’re coming to the main part:

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]):

This is what is executed first on the program start. We’re obtaining the free port that we’ll then use for our proxy. We start proxy in embedded mode and point some parameters to configure the run. Some of those parameters configure MITM module - they refer to files we have generated with the help of openssl tool.

We also load the plugin we have prepared. That plugin passes all the traffic through itself hence the headers are added to all the requests which go from your browser used for testing.

So, well, this is basically it. If you run the code you shouldn’t get any assertion exceptions. You can play with the expected messages to make sure your test successfully enters the page and hence bypasses authentication mechanism which means the header has been successfully added.

Hope the examples are clear and self-explanatory however if you still have the questions please send them to me using this form. I will amend the article according to your feedback.