From ff0e822b420fbbdd1a569961f4a5c8ca4b9638ec Mon Sep 17 00:00:00 2001
From: Will Starms <vilhelmen@gmail.com>
Date: Wed, 24 Jul 2024 21:31:53 -0500
Subject: [PATCH] Header-based auth class

---
 docs/user/authentication.rst | 21 +++++++++++++++++++
 src/requests/auth.py         | 11 ++++++++++
 tests/test_requests.py       | 39 +++++++++++++++++++++++++++++++++++-
 3 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/docs/user/authentication.rst b/docs/user/authentication.rst
index 0737bd319a..5d29bb6094 100644
--- a/docs/user/authentication.rst
+++ b/docs/user/authentication.rst
@@ -32,6 +32,27 @@ for using it::
 Providing the credentials in a tuple like this is exactly the same as the
 ``HTTPBasicAuth`` example above.
 
+Header Authentication
+--------------------
+
+Some services require authentication data in the header of the request.
+Multiple headers can be added to an authentication object to keep them separate
+from request data::
+
+    >>> from requests.auth import HTTPHeaderAuth
+    >>> auth = HTTPHeaderAuth({'Api-Key': '1234567890ABCDEF'}})
+    >>> response = requests.get('https://httpbin.org/headers', auth=auth)
+    >>> response.json()['headers']['Api-Key']
+    '1234567890abcdef'
+
+Be aware that keys in the authentication object will override headers set by
+the current request or session parameters::
+
+    >>> from requests.auth import HTTPHeaderAuth
+    >>> auth = HTTPHeaderAuth({'Api-Key': '1234567890ABCDEF'}})
+    >>> response = requests.get('https://httpbin.org/headers', headers={'Api-Key': '0000000000'}, auth=auth)
+    >>> response.json()['headers']['Api-Key']
+    '1234567890ABCDEF'
 
 netrc Authentication
 ~~~~~~~~~~~~~~~~~~~~
diff --git a/src/requests/auth.py b/src/requests/auth.py
index 4a7ce6dc14..5a00b63264 100644
--- a/src/requests/auth.py
+++ b/src/requests/auth.py
@@ -104,6 +104,17 @@ def __call__(self, r):
         return r
 
 
+class HTTPHeaderAuth(AuthBase):
+    """Attaches authentication headers to the given Request object."""
+
+    def __init__(self, headers):
+        self.headers = headers
+
+    def __call__(self, r):
+        r.headers.update(self.headers)
+        return r
+
+
 class HTTPDigestAuth(AuthBase):
     """Attaches HTTP Digest Authentication to the given Request object."""
 
diff --git a/tests/test_requests.py b/tests/test_requests.py
index b4e9fe92ae..bf347c6175 100644
--- a/tests/test_requests.py
+++ b/tests/test_requests.py
@@ -17,7 +17,7 @@
 
 import requests
 from requests.adapters import HTTPAdapter
-from requests.auth import HTTPDigestAuth, _basic_auth_str
+from requests.auth import HTTPDigestAuth, HTTPHeaderAuth, _basic_auth_str
 from requests.compat import (
     JSONDecodeError,
     Morsel,
@@ -703,6 +703,43 @@ def get_netrc_auth_mock(url):
         finally:
             requests.sessions.get_netrc_auth = old_auth
 
+    def test_header_auth(self, httpbin):
+        header_key = "Test-Header"
+        header_value = "1234567890ABCDEF"
+        auth = HTTPHeaderAuth({header_key: header_value})
+        url = httpbin("headers")
+
+        # Check header exists
+        r = requests.get(url, auth=auth)
+        assert r.json()["headers"][header_key] == header_value
+
+        # Make sure it's not a fluke
+        r = requests.get(url)
+        assert header_key not in r.json()["headers"]
+
+        # Verify auth header overrides provided headers
+        #  (not strictly a feature, but it's current behavior)
+        second_header_value = "NOT_RETURNED"
+        r = requests.get(url, headers={header_key: second_header_value}, auth=auth)
+        assert r.json()["headers"][header_key] == header_value
+
+        # Test with session
+        s = requests.session()
+        s.auth = auth
+        r = s.get(url)
+        assert r.json()["headers"][header_key] == header_value
+
+        # verify session header override
+        r = s.get(url, headers={header_key: second_header_value})
+        assert r.json()["headers"][header_key] == header_value
+
+        # Check sure multiple headers works
+        header_keys = ("Header-One", "Header-Two")
+        auth = HTTPHeaderAuth({key: key for key in header_keys})
+        r = requests.get(url, auth=auth)
+        returned_keys = r.json()["headers"].keys()
+        assert all(key in returned_keys for key in header_keys)
+
     def test_DIGEST_HTTP_200_OK_GET(self, httpbin):
         for authtype in self.digest_auth_algo:
             auth = HTTPDigestAuth("user", "pass")