httphound/example_tests/03-basic_header_regex.py

import re
from pathlib import Path
from httphound.main import BaseProxyTest, BackendConfig, ProxyConfig


class FailHeaderRegexTest(BaseProxyTest):
    """Simple test to demonstrate how to match for header regexes.

    This test should fail
    """
    def __init__(self):
        super().__init__()
        self.description = "Header regex match (Fail)"
        # The expected_header_patterns should contain
        # the header name and a pattern that is used to
        # match the received header value
        self.expected_header_patterns = {
            "x-test": r"\d{1}$",
        }

        self.proxy_config = ProxyConfig(
            binary_path=Path.home() / "bin/haproxy"
        )

        self.backend_config = BackendConfig(
            response_headers={"X-Test": "1234"},
        )

    async def run_test(self):
        await self.make_request()
        print(self.response_headers)
        assert "x-test" in self.response_headers
        assert re.match(r"\d{1}$", self.response_headers["x-test"])
        return True


class PassHeaderRegexTest(BaseProxyTest):
    """Simple test to demonstrate how to match for header regexes.
    """
    def __init__(self):
        super().__init__()
        self.description = "Header regex match (Pass)"
        # The expected_header_patterns should contain
        # the header name and a pattern that is used to
        # match the received header value
        self.expected_header_patterns = {
            "x-test": r"\d{4}$",
        }

        self.proxy_config = ProxyConfig(
            binary_path=Path.home() / "bin/haproxy"
        )

        self.backend_config = BackendConfig(
            response_headers={"X-Test": "1234"},
        )

    async def run_test(self):
        await self.make_request()
        return True