"""Tests for System User Protection in user_manager.py Tests cover: - get_request_user_id(): 1st defense layer - blocks System Users from HTTP headers - get_request_user_filepath(): 2nd defense layer - structural blocking via get_public_user_directory() - add_user(): 3rd defense layer - prevents creation of System User names - Defense layers integration tests """ import pytest from unittest.mock import MagicMock, patch import tempfile import folder_paths from app.user_manager import UserManager @pytest.fixture def mock_user_directory(): """Create a temporary user directory.""" with tempfile.TemporaryDirectory() as temp_dir: original_dir = folder_paths.get_user_directory() folder_paths.set_user_directory(temp_dir) yield temp_dir folder_paths.set_user_directory(original_dir) @pytest.fixture def user_manager(mock_user_directory): """Create a UserManager instance for testing.""" with patch('app.user_manager.args') as mock_args: mock_args.multi_user = True manager = UserManager() # Add a default user for testing manager.users = {"default": "default", "test_user_123": "Test User"} yield manager @pytest.fixture def mock_request(): """Create a mock request object.""" request = MagicMock() request.headers = {} return request class TestGetRequestUserId: """Tests for get_request_user_id() - 1st defense layer. Verifies: - System Users (__ prefix) in HTTP header are rejected with KeyError - Public Users pass through successfully """ def test_system_user_raises_error(self, user_manager, mock_request): """Test System User in header raises KeyError.""" mock_request.headers = {"comfy-user": "__system"} with patch('app.user_manager.args') as mock_args: mock_args.multi_user = True with pytest.raises(KeyError, match="Unknown user"): user_manager.get_request_user_id(mock_request) def test_system_user_cache_raises_error(self, user_manager, mock_request): """Test System User cache raises KeyError.""" mock_request.headers = {"comfy-user": "__cache"} with patch('app.user_manager.args') as mock_args: mock_args.multi_user = True with pytest.raises(KeyError, match="Unknown user"): user_manager.get_request_user_id(mock_request) def test_normal_user_works(self, user_manager, mock_request): """Test normal user access works.""" mock_request.headers = {"comfy-user": "default"} with patch('app.user_manager.args') as mock_args: mock_args.multi_user = True user_id = user_manager.get_request_user_id(mock_request) assert user_id == "default" def test_unknown_user_raises_error(self, user_manager, mock_request): """Test unknown user raises KeyError.""" mock_request.headers = {"comfy-user": "unknown_user"} with patch('app.user_manager.args') as mock_args: mock_args.multi_user = True with pytest.raises(KeyError, match="Unknown user"): user_manager.get_request_user_id(mock_request) class TestGetRequestUserFilepath: """Tests for get_request_user_filepath() - 2nd defense layer. Verifies: - Returns None when get_public_user_directory() returns None (System User) - Acts as backup defense if 1st layer is bypassed """ def test_system_user_returns_none(self, user_manager, mock_request, mock_user_directory): """Test System User returns None (structural blocking).""" # First, we need to mock get_request_user_id to return System User # But actually, get_request_user_id will raise KeyError first # So we test via get_public_user_directory returning None mock_request.headers = {"comfy-user": "default"} with patch('app.user_manager.args') as mock_args: mock_args.multi_user = True # Patch get_public_user_directory to return None for testing with patch.object(folder_paths, 'get_public_user_directory', return_value=None): result = user_manager.get_request_user_filepath(mock_request, "test.txt") assert result is None def test_normal_user_gets_path(self, user_manager, mock_request, mock_user_directory): """Test normal user gets valid filepath.""" mock_request.headers = {"comfy-user": "default"} with patch('app.user_manager.args') as mock_args: mock_args.multi_user = True path = user_manager.get_request_user_filepath(mock_request, "test.txt") assert path is not None assert "default" in path assert path.endswith("test.txt") class TestAddUser: """Tests for add_user() - 3rd defense layer (creation-time blocking). Verifies: - System User name (__ prefix) creation is rejected with ValueError - Sanitized usernames that become System User are also rejected """ def test_system_user_prefix_name_raises(self, user_manager): """Test System User prefix in name raises ValueError.""" with pytest.raises(ValueError, match="System User prefix not allowed"): user_manager.add_user("__system") def test_system_user_prefix_cache_raises(self, user_manager): """Test System User cache prefix raises ValueError.""" with pytest.raises(ValueError, match="System User prefix not allowed"): user_manager.add_user("__cache") def test_sanitized_system_user_prefix_raises(self, user_manager): """Test sanitized name becoming System User prefix raises ValueError (bypass prevention).""" # "__test" directly starts with System User prefix with pytest.raises(ValueError, match="System User prefix not allowed"): user_manager.add_user("__test") def test_normal_user_creation(self, user_manager, mock_user_directory): """Test normal user creation works.""" user_id = user_manager.add_user("Normal User") assert user_id is not None assert not user_id.startswith("__") assert "Normal-User" in user_id or "Normal_User" in user_id def test_empty_name_raises(self, user_manager): """Test empty name raises ValueError.""" with pytest.raises(ValueError, match="username not provided"): user_manager.add_user("") def test_whitespace_only_raises(self, user_manager): """Test whitespace-only name raises ValueError.""" with pytest.raises(ValueError, match="username not provided"): user_manager.add_user(" ") class TestDefenseLayers: """Integration tests for all three defense layers. Verifies: - Each defense layer blocks System Users independently - System User bypass is impossible through any layer """ def test_layer1_get_request_user_id(self, user_manager, mock_request): """Test 1st defense layer blocks System Users.""" mock_request.headers = {"comfy-user": "__system"} with patch('app.user_manager.args') as mock_args: mock_args.multi_user = True with pytest.raises(KeyError): user_manager.get_request_user_id(mock_request) def test_layer2_get_public_user_directory(self): """Test 2nd defense layer blocks System Users.""" result = folder_paths.get_public_user_directory("__system") assert result is None def test_layer3_add_user(self, user_manager): """Test 3rd defense layer blocks System User creation.""" with pytest.raises(ValueError): user_manager.add_user("__system")