From ebf7de81e7b080d60837d94f2110473adceccd53 Mon Sep 17 00:00:00 2001 From: Bilal Al Date: Thu, 21 Mar 2024 14:10:32 -0700 Subject: [PATCH 1/2] Added request decorator class --- splitio/api/request_decorator.py | 89 +++++++++++++++++++++++++++++ tests/api/test_request_decorator.py | 53 +++++++++++++++++ 2 files changed, 142 insertions(+) create mode 100644 splitio/api/request_decorator.py create mode 100644 tests/api/test_request_decorator.py diff --git a/splitio/api/request_decorator.py b/splitio/api/request_decorator.py new file mode 100644 index 00000000..fffb9bf5 --- /dev/null +++ b/splitio/api/request_decorator.py @@ -0,0 +1,89 @@ +"""Request Decorator module.""" + +import abc + +_FORBIDDEN_HEADERS = [ + "SplitSDKVersion", + "SplitMachineIp", + "SplitMachineName", + "SplitImpressionsMode", + "Host", + "Referrer", + "Content-Type", + "Content-Length", + "Content-Encoding", + "Accept", + "Keep-Alive", + "X-Fastly-Debug" +] + +class UserCustomHeaderDecorator(object, metaclass=abc.ABCMeta): + """User custom header decorator interface.""" + + @abc.abstractmethod + def get_header_overrides(self): + """ + Return a dictionary with all the user-defined custom headers. + + :return: Dictionary {String: String} + :rtype: Dict + """ + pass + +class NoOpHeaderDecorator(UserCustomHeaderDecorator): + """User custom header Class for no headers.""" + + def get_header_overrides(self): + """ + Return a dictionary with all the user-defined custom headers. + + :return: Dictionary {String: String} + :rtype: Dict + """ + return {} + +class RequestDecorator(object): + """Request decorator class for injecting User custom data.""" + + def __init__(self, user_custom_header_decorator=None): + """ + Class constructor. + + :param user_custom_header_decorator: User custom header decorator instance. + :type user_custom_header_decorator: splitio.api.request_decorator.UserCustomHeaderDecorator + """ + if user_custom_header_decorator is None: + user_custom_header_decorator = NoOpHeaderDecorator() + + self._user_custom_header_decorator = user_custom_header_decorator + + def decorate_headers(self, request_session): + """ + Use a passed header dictionary and append user custom headers from the UserCustomHeaderDecorator instance. + + :param request_session: HTTP Request session + :type request_session: requests.Session() + + :return: Updated Request session + :rtype: requests.Session() + """ + try: + custom_headers = self._user_custom_header_decorator.get_header_overrides() + for header in custom_headers: + if self._is_header_allowed(header): + request_session.headers[header] = custom_headers[header] + return request_session + except Exception as exc: + raise ValueError('Problem adding custom header in request decorator') from exc + + def _is_header_allowed(self, header): + """ + Verivy for a given header if it exists in the list of reserved forbidden headers + + :param header: Dictionary containing header + :type headers: Dict + + :return: True if does not exist in forbidden headers list, False otherwise + :rtype: Boolean + """ + return header not in _FORBIDDEN_HEADERS \ No newline at end of file diff --git a/tests/api/test_request_decorator.py b/tests/api/test_request_decorator.py new file mode 100644 index 00000000..ae04311c --- /dev/null +++ b/tests/api/test_request_decorator.py @@ -0,0 +1,53 @@ +"""Request Decorator test module.""" +import requests +import pytest + +from splitio.api.request_decorator import RequestDecorator, UserCustomHeaderDecorator, _FORBIDDEN_HEADERS + +class RequestDecoratorTests(object): + """Request Decorator test cases.""" + + def test_noop(self): + """Test no operation.""" + decorator = RequestDecorator() + session = requests.Session() + old_len = len(session.headers) + session = decorator.decorate_headers(session) + assert(len(session.headers) == old_len) + + def test_add_custom_headers(self): + """test adding custom headers.""" + + class MyCustomDecorator(UserCustomHeaderDecorator): + def get_header_overrides(self): + return {"UserCustomHeader": "value", "AnotherCustomHeader": "val"} + + decorator = RequestDecorator(MyCustomDecorator()) + session = requests.Session() + session = decorator.decorate_headers(session) + assert(session.headers["UserCustomHeader"] == "value") + assert(session.headers["AnotherCustomHeader"] == "val") + + def test_add_forbidden_headers(self): + """test adding forbidden headers.""" + + class MyCustomDecorator(UserCustomHeaderDecorator): + def get_header_overrides(self): + final_header = {"UserCustomHeader": "value"} + [final_header.update({header: "val"}) for header in _FORBIDDEN_HEADERS] + return final_header + + decorator = RequestDecorator(MyCustomDecorator()) + session = requests.Session() + session = decorator.decorate_headers(session) + assert(session.headers["UserCustomHeader"] == "value") + + def test_errors(self): + class MyCustomDecorator(UserCustomHeaderDecorator): + def get_header_overrides(self): + return ["MyCustomHeader"] + + decorator = RequestDecorator(MyCustomDecorator()) + session = requests.Session() + with pytest.raises(ValueError): + session = decorator.decorate_headers(session) From 03ffd39436cb8bf3ab3bf350d974dfb04dff9b8b Mon Sep 17 00:00:00 2001 From: Bilal Al Date: Fri, 22 Mar 2024 15:41:12 -0700 Subject: [PATCH 2/2] used lower case check for forbidden headers --- splitio/api/request_decorator.py | 2 +- tests/api/test_request_decorator.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/splitio/api/request_decorator.py b/splitio/api/request_decorator.py index fffb9bf5..417aab86 100644 --- a/splitio/api/request_decorator.py +++ b/splitio/api/request_decorator.py @@ -86,4 +86,4 @@ def _is_header_allowed(self, header): :return: True if does not exist in forbidden headers list, False otherwise :rtype: Boolean """ - return header not in _FORBIDDEN_HEADERS \ No newline at end of file + return header.lower() not in [forbidden.lower() for forbidden in _FORBIDDEN_HEADERS] \ No newline at end of file diff --git a/tests/api/test_request_decorator.py b/tests/api/test_request_decorator.py index ae04311c..0256b89b 100644 --- a/tests/api/test_request_decorator.py +++ b/tests/api/test_request_decorator.py @@ -35,6 +35,7 @@ class MyCustomDecorator(UserCustomHeaderDecorator): def get_header_overrides(self): final_header = {"UserCustomHeader": "value"} [final_header.update({header: "val"}) for header in _FORBIDDEN_HEADERS] + [final_header.update({header.lower(): "val"}) for header in _FORBIDDEN_HEADERS] return final_header decorator = RequestDecorator(MyCustomDecorator())