Add support to fetch from authenticated token lists

moisses89 · moisses89 · commit ea4f3c6ecf9b · 2025-12-18T15:06:58.000+01:00
diff --git a/safe_transaction_service/tokens/migrations/0015_tokenlist_auth_key_tokenlist_auth_type_and_more.py b/safe_transaction_service/tokens/migrations/0015_tokenlist_auth_key_tokenlist_auth_type_and_more.py
@@ -0,0 +1,28 @@
+# Generated by Django 5.2.5 on 2025-12-18 12:13
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('tokens', '0014_tokennotvalid'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='tokenlist',
+            name='auth_key',
+            field=models.CharField(blank=True, help_text='Header name or query param name (e.g. api_key, token)', max_length=100),
+        ),
+        migrations.AddField(
+            model_name='tokenlist',
+            name='auth_type',
+            field=models.CharField(choices=[('none', 'None'), ('api_key', 'API Key (Header)'), ('bearer', 'Bearer Token'), ('basic', 'Basic Auth'), ('query', 'Query Param')], default='none', max_length=20),
+        ),
+        migrations.AddField(
+            model_name='tokenlist',
+            name='auth_value',
+            field=models.TextField(blank=True, help_text='Token / API key / credentials'),
+        ),
+    ]
diff --git a/safe_transaction_service/tokens/models.py b/safe_transaction_service/tokens/models.py
@@ -2,7 +2,7 @@
 import os
 from json import JSONDecodeError
 from typing import Optional, TypedDict
-from urllib.parse import urljoin
+from urllib.parse import parse_qsl, urlencode, urljoin, urlparse, urlunparse
 
 from django.conf import settings
 from django.core.exceptions import ValidationError
@@ -341,12 +341,104 @@ class TokenList(models.Model):
     url = models.URLField(unique=True)
     description = models.CharField(max_length=200)
 
+    auth_type = models.CharField(
+        max_length=20,
+        choices=[
+            ("none", "None"),
+            ("api_key", "API Key (Header)"),
+            ("bearer", "Bearer Token"),
+            ("basic", "Basic Auth"),
+            ("query", "Query Param"),
+        ],
+        default="none",
+    )
+
+    auth_key = models.CharField(
+        max_length=100,
+        blank=True,
+        help_text=(
+            "Authentication key name. Usage depends on auth_type: "
+            "For 'api_key': header name (e.g., 'x-api-key'). "
+            "For 'bearer': header name (e.g., 'Authorization'). "
+            "For 'query': query parameter name (e.g., 'api_key'). "
+            "For 'basic': leave empty. "
+            "For 'none': leave empty."
+        ),
+    )
+
+    auth_value = models.TextField(
+        blank=True,
+        help_text=(
+            "Authentication credentials. Usage depends on auth_type: "
+            "For 'api_key': the API key value. "
+            "For 'bearer': the bearer token value (without 'Bearer ' prefix). "
+            "For 'basic': the full 'Basic <base64>' value (e.g., 'Basic dXNlcjpwYXNz'). "
+            "For 'query': the query parameter value. "
+            "For 'none': leave empty."
+        ),
+    )
+
     def __str__(self):
         return f"{self.description} token list"
 
+    def clean(self):
+        super().clean()
+
+        if self.auth_type == "none":
+            if self.auth_key or self.auth_value:
+                raise ValidationError(
+                    "auth_key and auth_value must be empty when auth_type is 'none'"
+                )
+
+        elif self.auth_type in {"bearer", "api_key", "query"}:
+            if not self.auth_key:
+                raise ValidationError(
+                    {"auth_key": "auth_key is required for this auth_type"}
+                )
+            if not self.auth_value:
+                raise ValidationError(
+                    {"auth_value": "auth_value is required for this auth_type"}
+                )
+
+        elif self.auth_type == "basic":
+            if not self.auth_value:
+                raise ValidationError(
+                    {"auth_value": "auth_value is required for basic auth"}
+                )
+
+    def _build_authenticated_request(self) -> tuple[str, dict]:
+        """
+        Builds the URL and headers with authentication based on the configured auth_type.
+        Supports bearer tokens, API keys (header-based), basic auth, and query parameters.
+
+        :return: Tuple of (authenticated_url, headers_dict)
+        """
+        url = self.url
+        headers: dict[str, str] = {}
+
+        if self.auth_type == "bearer":
+            header_name = self.auth_key if self.auth_key else "Authorization"
+            headers[header_name] = f"Bearer {self.auth_value}"
+
+        elif self.auth_type == "api_key":
+            headers[self.auth_key] = self.auth_value
+
+        elif self.auth_type == "basic":
+            headers["Authorization"] = self.auth_value
+
+        elif self.auth_type == "query":
+            parsed = urlparse(url)
+            query = dict(parse_qsl(parsed.query))
+            query[self.auth_key] = self.auth_value
+
+            url = urlunparse(parsed._replace(query=urlencode(query)))
+
+        return url, headers
+
     def get_tokens(self) -> list[TokenListToken]:
         try:
-            response = requests.get(self.url, timeout=5)
+            url, headers = self._build_authenticated_request()
+            response = requests.get(url, headers=headers, timeout=5)
             if response.ok:
                 tokens = response.json().get("tokens", [])
                 if not tokens:
diff --git a/safe_transaction_service/tokens/tests/test_models.py b/safe_transaction_service/tokens/tests/test_models.py
@@ -259,3 +259,102 @@ def test_get_tokens(self, requests_get: MagicMock):
 
         requests_get.return_value.json.return_value = token_list_mock
         self.assertEqual(len(token_list.get_tokens()), 35)
+
+    @mock.patch("requests.get")
+    def test_get_tokens_with_auth_none(self, requests_get: MagicMock):
+        token_list = TokenListFactory(
+            url="https://example.com/tokens.json",
+            auth_type="none",
+        )
+
+        requests_get.return_value.ok = True
+        requests_get.return_value.json.return_value = token_list_mock
+
+        tokens = token_list.get_tokens()
+
+        self.assertEqual(len(tokens), 35)
+        requests_get.assert_called_once_with(
+            "https://example.com/tokens.json", headers={}, timeout=5
+        )
+
+    @mock.patch("requests.get")
+    def test_get_tokens_with_auth_bearer(self, requests_get: MagicMock):
+        token_list = TokenListFactory(
+            url="https://example.com/tokens.json",
+            auth_type="bearer",
+            auth_value="my-secret-token",
+        )
+
+        requests_get.return_value.ok = True
+        requests_get.return_value.json.return_value = token_list_mock
+
+        tokens = token_list.get_tokens()
+
+        self.assertEqual(len(tokens), 35)
+        requests_get.assert_called_once_with(
+            "https://example.com/tokens.json",
+            headers={"Authorization": "Bearer my-secret-token"},
+            timeout=5,
+        )
+
+    @mock.patch("requests.get")
+    def test_get_tokens_with_auth_api_key(self, requests_get: MagicMock):
+        token_list = TokenListFactory(
+            url="https://example.com/tokens.json",
+            auth_type="api_key",
+            auth_key="X-API-Key",
+            auth_value="my-api-key-value",
+        )
+
+        requests_get.return_value.ok = True
+        requests_get.return_value.json.return_value = token_list_mock
+
+        tokens = token_list.get_tokens()
+
+        self.assertEqual(len(tokens), 35)
+        requests_get.assert_called_once_with(
+            "https://example.com/tokens.json",
+            headers={"X-API-Key": "my-api-key-value"},
+            timeout=5,
+        )
+
+    @mock.patch("requests.get")
+    def test_get_tokens_with_auth_basic(self, requests_get: MagicMock):
+        token_list = TokenListFactory(
+            url="https://example.com/tokens.json",
+            auth_type="basic",
+            auth_value="Basic dXNlcm5hbWU6cGFzc3dvcmQ=",
+        )
+
+        requests_get.return_value.ok = True
+        requests_get.return_value.json.return_value = token_list_mock
+
+        tokens = token_list.get_tokens()
+
+        self.assertEqual(len(tokens), 35)
+        requests_get.assert_called_once_with(
+            "https://example.com/tokens.json",
+            headers={"Authorization": "Basic dXNlcm5hbWU6cGFzc3dvcmQ="},
+            timeout=5,
+        )
+
+    @mock.patch("requests.get")
+    def test_get_tokens_with_auth_query(self, requests_get: MagicMock):
+        token_list = TokenListFactory(
+            url="https://example.com/tokens.json",
+            auth_type="query",
+            auth_key="api_key",
+            auth_value="my-query-param-key",
+        )
+
+        requests_get.return_value.ok = True
+        requests_get.return_value.json.return_value = token_list_mock
+
+        tokens = token_list.get_tokens()
+
+        self.assertEqual(len(tokens), 35)
+        requests_get.assert_called_once_with(
+            "https://example.com/tokens.json?api_key=my-query-param-key",
+            headers={},
+            timeout=5,
+        )
