fix: add missing uv index env var for auth (#751)

yeisonvargasf · web-flow · commit 1136c27e6ee3 · 2025-05-13T12:00:33.000-05:00
diff --git a/safety/tool/auth.py b/safety/tool/auth.py
@@ -0,0 +1,35 @@
+import base64
+import json
+
+import typer
+
+
+def index_credentials(ctx: typer.Context) -> str:
+    """
+    Returns the index credentials for the current context.
+    This should be used together with user:index_credential for index
+    basic auth.
+
+    Args:
+        ctx (typer.Context): The context.
+
+    Returns:
+        str: The index credentials.
+    """
+    api_key = None
+    token = None
+
+    if auth := getattr(ctx.obj, "auth", None):
+        client = auth.client
+        token = client.token.get("access_token") if client.token else None
+        api_key = client.api_key
+
+    auth_envelop = json.dumps(
+        {
+            "version": "1.0",
+            "access_token": token,
+            "api_key": api_key,
+            "project_id": ctx.obj.project.id if ctx.obj.project else None,
+        }
+    )
+    return base64.urlsafe_b64encode(auth_envelop.encode("utf-8")).decode("utf-8")
diff --git a/safety/tool/pip/main.py b/safety/tool/pip/main.py
@@ -1,5 +1,3 @@
-import base64
-import json
 import logging
 import re
 import shutil
@@ -19,6 +17,7 @@
 from safety.tool.resolver import get_unwrapped_command
 
 from safety.console import main_console
+from safety.tool.auth import index_credentials
 
 
 logger = logging.getLogger(__name__)
@@ -137,26 +136,6 @@ def reset_system(cls, console: Console = main_console):
         except Exception:
             console.print("Failed to reset PIP global settings.")
 
-    @classmethod
-    def index_credentials(cls, ctx: typer.Context):
-        api_key = None
-        token = None
-
-        if auth := getattr(ctx.obj, "auth", None):
-            client = auth.client
-            token = client.token.get("access_token") if client.token else None
-            api_key = client.api_key
-
-        auth_envelop = json.dumps(
-            {
-                "version": "1.0",
-                "access_token": token,
-                "api_key": api_key,
-                "project_id": ctx.obj.project.id if ctx.obj.project else None,
-            }
-        )
-        return base64.urlsafe_b64encode(auth_envelop.encode("utf-8")).decode("utf-8")
-
     @classmethod
     def default_index_url(cls) -> str:
         return "https://pypi.org/simple/"
@@ -168,7 +147,7 @@ def build_index_url(cls, ctx: typer.Context, index_url: Optional[str]) -> str:
 
         url = urlsplit(index_url)
 
-        encoded_auth = cls.index_credentials(ctx)
+        encoded_auth = index_credentials(ctx)
         netloc = f"user:{encoded_auth}@{url.netloc}"
 
         if type(url.netloc) is bytes:
diff --git a/safety/tool/uv/command.py b/safety/tool/uv/command.py
@@ -1,6 +1,9 @@
 from typing import List
+
+import typer
 from safety.tool.intents import ToolIntentionType
 from safety.tool.pip.parser import PipParser
+from safety.tool.auth import index_credentials
 from ..pip.command import PipCommand, PipInstallCommand, PipGenericCommand
 from safety_schemas.models.events.types import ToolType
 
@@ -35,6 +38,18 @@ def should_track_state(self) -> bool:
 
         return any(cmd in command_str for cmd in package_modifying_commands)
 
+    def env(self, ctx: typer.Context) -> dict:
+        env = super().env(ctx)
+
+        env.update(
+            {
+                "UV_INDEX_SAFETY_USERNAME": "user",
+                "UV_INDEX_SAFETY_PASSWORD": index_credentials(ctx),
+            }
+        )
+
+        return env
+
     @classmethod
     def from_args(cls, args: List[str], **kwargs):
         pip_parser = PipParser()
diff --git a/tests/tool/test_auth.py b/tests/tool/test_auth.py
@@ -0,0 +1,147 @@
+# type: ignore
+import unittest
+import json
+import base64
+from unittest.mock import MagicMock
+
+import typer
+from safety.tool.auth import index_credentials
+
+
+class TestIndexCredentials(unittest.TestCase):
+    """
+    Test cases for index_credentials function.
+    """
+
+    def test_index_credentials_with_full_auth_object(self):
+        """
+        Test index_credentials when ctx.obj.auth is fully populated with token and api_key.
+        """
+        ctx = MagicMock(spec=typer.Context)
+        ctx.obj = MagicMock()
+        ctx.obj.auth.client.token = {"access_token": "test_token"}
+        ctx.obj.auth.client.api_key = "test_api_key"
+        ctx.obj.project.id = "test_project_id"
+
+        result = index_credentials(ctx)
+
+        decoded = json.loads(
+            base64.urlsafe_b64decode(result.encode("utf-8")).decode("utf-8")
+        )
+
+        self.assertEqual(decoded["version"], "1.0")
+        self.assertEqual(decoded["access_token"], "test_token")
+        self.assertEqual(decoded["api_key"], "test_api_key")
+        self.assertEqual(decoded["project_id"], "test_project_id")
+
+    def test_index_credentials_with_missing_token(self):
+        """
+        Test index_credentials when token is None but api_key is present.
+        """
+        ctx = MagicMock(spec=typer.Context)
+        ctx.obj = MagicMock()
+        ctx.obj.auth.client.token = None
+        ctx.obj.auth.client.api_key = "test_api_key"
+        ctx.obj.project.id = "test_project_id"
+
+        result = index_credentials(ctx)
+
+        decoded = json.loads(
+            base64.urlsafe_b64decode(result.encode("utf-8")).decode("utf-8")
+        )
+        self.assertEqual(decoded["version"], "1.0")
+        self.assertIsNone(decoded["access_token"])
+        self.assertEqual(decoded["api_key"], "test_api_key")
+        self.assertEqual(decoded["project_id"], "test_project_id")
+
+    def test_index_credentials_with_missing_api_key(self):
+        """
+        Test index_credentials when api_key is None but token is present.
+        """
+        ctx = MagicMock(spec=typer.Context)
+        ctx.obj = MagicMock()
+        ctx.obj.auth.client.token = {"access_token": "test_token"}
+        ctx.obj.auth.client.api_key = None
+        ctx.obj.project.id = "test_project_id"
+
+        result = index_credentials(ctx)
+
+        decoded = json.loads(
+            base64.urlsafe_b64decode(result.encode("utf-8")).decode("utf-8")
+        )
+
+        self.assertEqual(decoded["version"], "1.0")
+        self.assertEqual(decoded["access_token"], "test_token")
+        self.assertIsNone(decoded["api_key"])
+        self.assertEqual(decoded["project_id"], "test_project_id")
+
+    def test_index_credentials_with_no_auth(self):
+        """
+        Test index_credentials when ctx.obj.auth is None.
+        """
+
+        ctx = MagicMock(spec=typer.Context)
+        ctx.obj = MagicMock()
+        ctx.obj.auth = None
+        ctx.obj.project.id = "test_project_id"
+
+        result = index_credentials(ctx)
+
+        decoded = json.loads(
+            base64.urlsafe_b64decode(result.encode("utf-8")).decode("utf-8")
+        )
+
+        self.assertEqual(decoded["version"], "1.0")
+        self.assertIsNone(decoded["access_token"])
+        self.assertIsNone(decoded["api_key"])
+        self.assertEqual(decoded["project_id"], "test_project_id")
+
+    def test_index_credentials_with_no_project(self):
+        """
+        Test index_credentials when ctx.obj.project is None.
+        """
+
+        ctx = MagicMock(spec=typer.Context)
+        ctx.obj = MagicMock()
+        ctx.obj.auth.client.token = {"access_token": "test_token"}
+        ctx.obj.auth.client.api_key = "test_api_key"
+        ctx.obj.project = None
+
+        result = index_credentials(ctx)
+
+        decoded = json.loads(
+            base64.urlsafe_b64decode(result.encode("utf-8")).decode("utf-8")
+        )
+
+        self.assertEqual(decoded["version"], "1.0")
+        self.assertEqual(decoded["access_token"], "test_token")
+        self.assertEqual(decoded["api_key"], "test_api_key")
+        self.assertIsNone(decoded["project_id"])
+
+    def test_index_credentials_correct_encoding(self):
+        """
+        Test that index_credentials correctly encodes the credentials in base64url format.
+        """
+
+        ctx = MagicMock(spec=typer.Context)
+        ctx.obj = MagicMock()
+        ctx.obj.auth.client.token = {"access_token": "test_token"}
+        ctx.obj.auth.client.api_key = "test_api_key"
+        ctx.obj.project.id = "test_project_id"
+
+        result = index_credentials(ctx)
+
+        expected_json = json.dumps(
+            {
+                "version": "1.0",
+                "access_token": "test_token",
+                "api_key": "test_api_key",
+                "project_id": "test_project_id",
+            }
+        )
+
+        expected_encoded = base64.urlsafe_b64encode(
+            expected_json.encode("utf-8")
+        ).decode("utf-8")
+
+        self.assertEqual(result, expected_encoded)
diff --git a/tests/tool/uv/test_uv_command.py b/tests/tool/uv/test_uv_command.py
@@ -0,0 +1,89 @@
+# type: ignore
+import unittest
+from unittest.mock import patch, MagicMock
+
+import typer
+from safety.tool.uv.command import UvCommand
+
+
+class TestUvCommand(unittest.TestCase):
+    """
+    Test cases for UvCommand functionality.
+    """
+
+    def setUp(self):
+        """
+        Set up test environment before each test method.
+        """
+        self.command = UvCommand(["uv", "pip", "install", "package"])
+
+    def test_env_preserves_existing_variables(self):
+        """
+        Test that env() method does not replace existing environment variables.
+        """
+        ctx = MagicMock(spec=typer.Context)
+
+        existing_env = {
+            "EXISTING_VAR": "existing_value",
+            "ANOTHER_VAR": "another_value",
+        }
+
+        with patch(
+            "safety.tool.pip.command.PipCommand.env", return_value=existing_env
+        ) as mock_super_env:
+            with patch(
+                "safety.tool.uv.command.index_credentials",
+                return_value="mock_credentials",
+            ):
+                result_env = self.command.env(ctx)
+
+                self.assertEqual(result_env["EXISTING_VAR"], "existing_value")
+                self.assertEqual(result_env["ANOTHER_VAR"], "another_value")
+                mock_super_env.assert_called_once_with(ctx)
+
+    def test_env_adds_uv_credentials_variables(self):
+        """
+        Test that env() method always adds UV_INDEX_SAFETY_USERNAME and
+        UV_INDEX_SAFETY_PASSWORD environment variables.
+        """
+        ctx = MagicMock(spec=typer.Context)
+
+        mock_credentials = "mock_credentials_value"
+
+        with patch(
+            "safety.tool.uv.command.index_credentials", return_value=mock_credentials
+        ):
+            with patch("safety.tool.pip.command.PipCommand.env", return_value={}):
+                result_env = self.command.env(ctx)
+
+                self.assertIn("UV_INDEX_SAFETY_USERNAME", result_env)
+                self.assertIn("UV_INDEX_SAFETY_PASSWORD", result_env)
+                self.assertEqual(result_env["UV_INDEX_SAFETY_USERNAME"], "user")
+                self.assertEqual(
+                    result_env["UV_INDEX_SAFETY_PASSWORD"], mock_credentials
+                )
+
+    def test_env_combines_parent_env_with_uv_credentials(self):
+        """
+        Test that env() method properly combines parent environment with UV credentials.
+        """
+        ctx = MagicMock(spec=typer.Context)
+
+        existing_env = {"EXISTING_VAR": "existing_value", "PATH": "/usr/bin:/bin"}
+
+        mock_credentials = "mock_credentials_value"
+
+        with patch(
+            "safety.tool.uv.command.index_credentials", return_value=mock_credentials
+        ):
+            with patch(
+                "safety.tool.pip.command.PipCommand.env", return_value=existing_env
+            ):
+                result_env = self.command.env(ctx)
+
+                self.assertEqual(result_env["EXISTING_VAR"], "existing_value")
+                self.assertEqual(result_env["PATH"], "/usr/bin:/bin")
+                self.assertEqual(result_env["UV_INDEX_SAFETY_USERNAME"], "user")
+                self.assertEqual(
+                    result_env["UV_INDEX_SAFETY_PASSWORD"], mock_credentials
+                )
