From eb40c67385ddbde0651d29c5b3786b0206e7220e Mon Sep 17 00:00:00 2001
From: hifumi <supersonik@gmx.de>
Date: Tue, 4 Aug 2026 16:20:24 +0200
Subject: [PATCH] feat(vaults): support custom headers on API vault

---
 docs/guide/vaults.md    |  4 ++++
 unshackle/vaults/API.py | 44 +++++++++++++++++++++++++++++++----------
 2 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/docs/guide/vaults.md b/docs/guide/vaults.md
index f62926b8..eedeca51 100644
--- a/docs/guide/vaults.md
+++ b/docs/guide/vaults.md
@@ -208,6 +208,9 @@ key_vaults:
     name: myapi
     uri: https://vault.example.com/api
     token: "your-bearer-token"
+    headers:
+      CF-Access-Client-Id: "CLIENT_ID_HERE"
+      CF-Access-Client-Secret: "CLIENT_SECRET_HERE"
 ```
 
 | Field | Required | Description |
@@ -218,6 +221,7 @@ key_vaults:
 | `token` | yes | Bearer token, sent as `Authorization: Bearer <token>`. |
 | `no_push` | no | If `true`, keys are read but never written. |
 | `timeout` | no | Request timeout in seconds; defaults to `vault_timeout`. |
+| `headers` | no | Extra HTTP `headers` sent with every request, for example the Cloudflare Access service-token headers `CF-Access-Client-Id` and `CF-Access-Client-Secret`. |
 
 The backend talks to endpoints under `uri` (for example `GET {uri}/{service}/{kid}` to look up a single key) and interprets a numeric `code` field in each JSON response to detect errors such as an invalid token, rate limiting, or an invalid service tag.
 
diff --git a/unshackle/vaults/API.py b/unshackle/vaults/API.py
index 257313d3..25dd73b8 100644
--- a/unshackle/vaults/API.py
+++ b/unshackle/vaults/API.py
@@ -10,7 +10,7 @@ from unshackle.core.vault import Vault
 class API(Vault):
     """Key Vault using a simple RESTful HTTP API call."""
 
-    def __init__(self, name: str, uri: str, token: str, no_push: bool = False, timeout: float = 10.0):
+    def __init__(self, name: str, uri: str, token: str, no_push: bool = False, headers: Optional[dict[str, str]] = None, timeout: float = 10.0):
         super().__init__(name, no_push)
         self.uri = uri.rstrip("/")
         self.timeout = timeout
@@ -18,13 +18,21 @@ class API(Vault):
         self.session.headers.update({"User-Agent": f"unshackle v{__version__}"})
         self.session.headers.update({"Authorization": f"Bearer {token}"})
 
+        if headers:
+            self.session.headers.update(headers)
+
     def get_key(self, kid: Union[UUID, str], service: str) -> Optional[str]:
         if isinstance(kid, UUID):
             kid = kid.hex
 
-        data = self.session.get(
-            url=f"{self.uri}/{service.lower()}/{kid}", headers={"Accept": "application/json"}, timeout=self.timeout
-        ).json()
+        response = self.session.get(
+            url=f"{self.uri}/{service.lower()}/{kid}",
+            headers={"Accept": "application/json"},
+            timeout=self.timeout,
+            allow_redirects=False
+        )
+        response.raise_for_status()
+        data = response.json()
 
         code = int(data.get("code", 0))
         message = data.get("message")
@@ -52,12 +60,15 @@ class API(Vault):
         page = 1
 
         while True:
-            data = self.session.get(
+            response = self.session.get(
                 url=f"{self.uri}/{service.lower()}",
                 params={"page": page, "total": 10},
                 headers={"Accept": "application/json"},
                 timeout=self.timeout,
-            ).json()
+                allow_redirects=False
+            )
+            response.raise_for_status()
+            data = response.json()
 
             code = int(data.get("code", 0))
             message = data.get("message")
@@ -90,12 +101,15 @@ class API(Vault):
         if isinstance(kid, UUID):
             kid = kid.hex
 
-        data = self.session.post(
+        response = self.session.post(
             url=f"{self.uri}/{service.lower()}/{kid}",
             json={"content_key": key},
             headers={"Accept": "application/json"},
             timeout=self.timeout,
-        ).json()
+            allow_redirects=False
+        )
+        response.raise_for_status()
+        data = response.json()
 
         code = int(data.get("code", 0))
         message = data.get("message")
@@ -141,6 +155,7 @@ class API(Vault):
                     json={"content_keys": batch_keys},
                     headers={"Accept": "application/json"},
                     timeout=self.timeout,
+                    allow_redirects=False
                 )
 
                 # Check for HTTP errors that suggest batch is too large
@@ -150,7 +165,8 @@ class API(Vault):
                     else:
                         batch_size = 1
                     continue
-
+                
+                response.raise_for_status()
                 data = response.json()
             except Exception:
                 # JSON decode error or connection issue - try smaller batch
@@ -187,7 +203,14 @@ class API(Vault):
         return total_added
 
     def get_services(self) -> Iterator[str]:
-        data = self.session.post(url=self.uri, headers={"Accept": "application/json"}, timeout=self.timeout).json()
+        response = self.session.post(
+            url=self.uri,
+            headers={"Accept": "application/json"},
+            timeout=self.timeout,
+            allow_redirects=False
+        )
+        response.raise_for_status()
+        data = response.json()
 
         code = int(data.get("code", 0))
         message = data.get("message")
@@ -227,3 +250,4 @@ class Exceptions:
 
     class ContentKeyInvalid(Exception):
         """The Content Key is invalid."""
+        
\ No newline at end of file