diff --git a/libs/labelbox/src/labelbox/client.py b/libs/labelbox/src/labelbox/client.py index 082d5db01..75755afcc 100644 --- a/libs/labelbox/src/labelbox/client.py +++ b/libs/labelbox/src/labelbox/client.py @@ -8,6 +8,7 @@ import urllib.parse import warnings from collections import defaultdict +from collections.abc import Sequence from types import MappingProxyType from typing import Any, Callable, Dict, List, Optional, Set, Union @@ -31,6 +32,7 @@ from labelbox.pagination import PaginatedCollection from labelbox.project_validation import _CoreProjectInput from labelbox.schema import role +from labelbox.schema.api_key import ApiKey, IpAddressOrNetwork from labelbox.schema.catalog import Catalog from labelbox.schema.data_row import DataRow from labelbox.schema.data_row_metadata import DataRowMetadataOntology @@ -78,10 +80,9 @@ ) from labelbox.schema.slice import CatalogSlice, ModelSlice from labelbox.schema.task import DataUpsertTask, Task -from labelbox.schema.user import User from labelbox.schema.taskstatus import TaskStatus -from labelbox.schema.api_key import ApiKey from labelbox.schema.timeunit import TimeUnit +from labelbox.schema.user import User logger = logging.getLogger(__name__) @@ -2446,6 +2447,7 @@ def create_api_key( validity: int = 0, time_unit: TimeUnit = TimeUnit.SECOND, refresh_cache: bool = False, + allowed_ip_cidrs: Optional[Sequence[IpAddressOrNetwork]] = None, ) -> Dict[str, str]: """Creates a new API key. @@ -2456,6 +2458,9 @@ def create_api_key( validity (int, optional): The validity period of the API key. Defaults to 0 (no expiration). time_unit (TimeUnit, optional): The time unit for the validity period. Defaults to TimeUnit.SECOND. refresh_cache (bool, optional): Whether to refresh cached permissions and roles. Defaults to False. + allowed_ip_cidrs: Optional sequence of IP addresses or CIDR ranges + allowed to use the key. Entries may be strings or standard + :mod:`ipaddress` address/network objects. Returns: Dict[str, str]: A dictionary containing the created API key information. @@ -2471,7 +2476,13 @@ def create_api_key( delattr(self, "_cached_available_api_key_roles") return ApiKey.create_api_key( - self, name, user, role, validity, time_unit + self, + name, + user, + role, + validity, + time_unit, + allowed_ip_cidrs, ) def get_api_keys(self, include_expired: bool = False) -> List[ApiKey]: diff --git a/libs/labelbox/src/labelbox/schema/api_key.py b/libs/labelbox/src/labelbox/schema/api_key.py index f15964871..bd7323b6a 100644 --- a/libs/labelbox/src/labelbox/schema/api_key.py +++ b/libs/labelbox/src/labelbox/schema/api_key.py @@ -1,20 +1,34 @@ -from datetime import datetime, timezone -from typing import Any, Dict, List, Optional, Union, TYPE_CHECKING import logging +from collections.abc import Sequence +from datetime import datetime, timezone +from ipaddress import ( + IPv4Address, + IPv4Network, + IPv6Address, + IPv6Network, + ip_address, + ip_network, +) +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union + +from lbox.exceptions import LabelboxError from labelbox.orm.db_object import DbObject from labelbox.orm.model import Field +from labelbox.schema.role import Role, format_role from labelbox.schema.timeunit import TimeUnit -from lbox.exceptions import LabelboxError - from labelbox.schema.user import User -from labelbox.schema.role import Role, format_role if TYPE_CHECKING: from labelbox import Client logger = logging.getLogger(__name__) +MAX_API_KEY_ALLOWED_IP_CIDRS = 50 +IpAddressOrNetwork = Union[ + str, IPv4Address, IPv6Address, IPv4Network, IPv6Network +] + class ApiKey(DbObject): """Represents an API key in the Labelbox system. @@ -30,6 +44,8 @@ class ApiKey(DbObject): expires_at_epoch (int): Expiration time as Unix timestamp created_by_user (str): ID of the user who created this API key user (str): ID of the user this API key belongs to + allowed_ip_cidrs (List[str]): IP addresses and CIDR ranges allowed to + use this API key """ name = Field.String("name") @@ -39,6 +55,9 @@ class ApiKey(DbObject): expires_at_epoch = Field.Int("expires_at_epoch") created_by_user_id = Field.String("created_by_user_id") user_id = Field.String("user_id") + allowed_ip_cidrs = Field.List( + str, graphql_type="String", name="allowed_ip_cidrs" + ) @property def created_by(self) -> Optional["User"]: @@ -152,6 +171,58 @@ def _get_current_user_permissions(client: "Client") -> List[str]: return perms + @staticmethod + def _validate_allowed_ip_cidrs( + allowed_ip_cidrs: Optional[Sequence[IpAddressOrNetwork]], + ) -> Optional[List[str]]: + """Validate and normalize an API key IP allowlist.""" + if allowed_ip_cidrs is None: + return None + if isinstance(allowed_ip_cidrs, (str, bytes)) or not isinstance( + allowed_ip_cidrs, Sequence + ): + raise ValueError( + "allowed_ip_cidrs must be a sequence of IP addresses or " + "CIDR ranges" + ) + normalized: List[str] = [] + for value in allowed_ip_cidrs: + if isinstance(value, (IPv4Address, IPv6Address)): + entry = str(value) + elif isinstance(value, (IPv4Network, IPv6Network)): + entry = str(value) + elif isinstance(value, str): + entry = value.strip() + if not entry: + raise ValueError( + "allowed_ip_cidrs entries must not be empty" + ) + try: + entry = str(ip_address(entry)) + except ValueError: + try: + entry = str(ip_network(entry, strict=True)) + except ValueError as error: + raise ValueError( + f"Invalid IP address or CIDR range: {value!r}" + ) from error + else: + raise ValueError( + "allowed_ip_cidrs entries must be strings or ipaddress " + "address/network objects" + ) + + if entry not in normalized: + normalized.append(entry) + + if len(normalized) > MAX_API_KEY_ALLOWED_IP_CIDRS: + raise ValueError( + "allowed_ip_cidrs cannot contain more than " + f"{MAX_API_KEY_ALLOWED_IP_CIDRS} unique entries" + ) + + return normalized + @staticmethod def get_api_keys( client: "Client", include_expired: bool = False @@ -180,6 +251,7 @@ def get_api_keys( userId userEmail createdByUserId + allowedIpCidrs } apiKeysOtherUsers { id @@ -191,6 +263,7 @@ def get_api_keys( expiresAtEpoch createdByUserId userEmail + allowedIpCidrs } } } @@ -205,13 +278,25 @@ def get_api_keys( current_time = datetime.now(timezone.utc) for key_data in response["user"].get("apiKeys", []): - api_key = ApiKey(client, key_data) + api_key = ApiKey( + client, + { + **key_data, + "allowedIpCidrs": key_data.get("allowedIpCidrs") or [], + }, + ) # Only add if we want to include expired keys OR if the key has not expired. if include_expired or api_key.expired_at > current_time: all_keys.append(api_key) for key_data in response["user"].get("apiKeysOtherUsers", []): - api_key = ApiKey(client, key_data) + api_key = ApiKey( + client, + { + **key_data, + "allowedIpCidrs": key_data.get("allowedIpCidrs") or [], + }, + ) if include_expired or api_key.expired_at > current_time: all_keys.append(api_key) @@ -297,6 +382,7 @@ def create_api_key( role: Union["Role", str], validity: int = 0, time_unit: TimeUnit = TimeUnit.SECOND, + allowed_ip_cidrs: Optional[Sequence[IpAddressOrNetwork]] = None, ) -> Dict[str, str]: """Creates a new API key using the provided client. @@ -307,6 +393,10 @@ def create_api_key( role (Union[Role, str]): Permission role for the API key (Role enum or string) validity (int, optional): Validity period value (must be positive). Defaults to 0. time_unit (TimeUnit, optional): Time unit for validity period. Defaults to TimeUnit.SECOND. + allowed_ip_cidrs: Optional sequence of IP addresses or CIDR ranges + allowed to use the key. Entries may be strings or standard + :mod:`ipaddress` address/network objects. A maximum of 50 + unique entries is supported. Returns: Dict[str, str]: Dictionary containing the created API key details including id and jwt @@ -318,6 +408,10 @@ def create_api_key( if not name or not isinstance(name, str): raise ValueError("name must be a non-empty string") + normalized_allowed_ip_cidrs = ApiKey._validate_allowed_ip_cidrs( + allowed_ip_cidrs + ) + user_email = user.email if hasattr(user, "email") else user if not user_email or not isinstance(user_email, str): raise ValueError("user must be a User object or a valid email") @@ -368,9 +462,9 @@ def create_api_key( ) query_str = """ - mutation CreateUserApiKeyPyApi($name: String!, $userEmail: String!, $role: String, $validitySeconds: Int) { + mutation CreateUserApiKeyPyApi($name: String!, $userEmail: String!, $role: String, $validitySeconds: Int, $allowedIpCidrs: [String!]) { createApiKey( - data: { name: $name, targetUserEmailId: $userEmail, role: $role, validitySeconds: $validitySeconds } + data: { name: $name, targetUserEmailId: $userEmail, role: $role, validitySeconds: $validitySeconds, allowedIpCidrs: $allowedIpCidrs } ) { id jwt @@ -383,6 +477,7 @@ def create_api_key( "userEmail": user_email, "role": server_role_name, "validitySeconds": validity_seconds, + "allowedIpCidrs": normalized_allowed_ip_cidrs, } try: diff --git a/libs/labelbox/tests/unit/schema/test_api_key.py b/libs/labelbox/tests/unit/schema/test_api_key.py new file mode 100644 index 000000000..3ed967594 --- /dev/null +++ b/libs/labelbox/tests/unit/schema/test_api_key.py @@ -0,0 +1,151 @@ +from datetime import datetime, timezone +from ipaddress import ip_address, ip_network +from unittest.mock import MagicMock, patch + +import pytest + +from labelbox.schema.api_key import ApiKey +from labelbox.schema.timeunit import TimeUnit + + +def _api_key_data(**overrides): + data = { + "id": "key-1", + "name": "Restricted key", + "createdAt": "2026-09-17T12:00:00.000Z", + "updatedAt": "2026-09-17T12:00:00.000Z", + "revoked": False, + "expiresAtEpoch": int( + datetime(2030, 1, 1, tzinfo=timezone.utc).timestamp() + ), + "createdByUserId": "user-1", + "userId": "user-1", + "allowedIpCidrs": ["203.0.113.10", "198.51.100.0/24"], + } + data.update(overrides) + return data + + +def test_create_api_key_validates_and_serializes_ip_allowlist(): + client = MagicMock() + client.execute.return_value = { + "createApiKey": {"id": "key-1", "jwt": "secret"} + } + + with ( + patch.object(ApiKey, "_get_user", return_value="user-1"), + patch.object( + ApiKey, "_get_available_api_key_roles", return_value=["Admin"] + ), + ): + result = ApiKey.create_api_key( + client, + name="Restricted key", + user="person@example.com", + role="Admin", + validity=5, + time_unit=TimeUnit.MINUTE, + allowed_ip_cidrs=[ + "203.0.113.10", + ip_network("198.51.100.0/24"), + ip_address("2001:0db8::1"), + "203.0.113.10", + ], + ) + + assert result == {"id": "key-1", "jwt": "secret"} + query, variables = client.execute.call_args.args + assert "$allowedIpCidrs: [String!]" in query + assert variables["allowedIpCidrs"] == [ + "203.0.113.10", + "198.51.100.0/24", + "2001:db8::1", + ] + + +@pytest.mark.parametrize( + "allowed_ip_cidrs", + [ + "203.0.113.10", + ["not-an-ip"], + ["192.0.2.10/24"], + [""], + [1234], + [f"192.0.2.{index}" for index in range(51)], + ], +) +def test_create_api_key_rejects_invalid_ip_allowlist_before_api_calls( + allowed_ip_cidrs, +): + client = MagicMock() + + with pytest.raises(ValueError, match="allowed_ip_cidrs|Invalid IP"): + ApiKey.create_api_key( + client, + name="Restricted key", + user="person@example.com", + role="Admin", + validity=5, + time_unit=TimeUnit.MINUTE, + allowed_ip_cidrs=allowed_ip_cidrs, + ) + + client.execute.assert_not_called() + client.get_organization.assert_not_called() + + +def test_get_api_keys_exposes_allowed_ip_cidrs(): + client = MagicMock() + client.execute.return_value = { + "user": { + "apiKeys": [_api_key_data()], + "apiKeysOtherUsers": [], + } + } + + api_keys = ApiKey.get_api_keys(client) + + assert len(api_keys) == 1 + assert api_keys[0].allowed_ip_cidrs == [ + "203.0.113.10", + "198.51.100.0/24", + ] + assert "allowedIpCidrs" in client.execute.call_args.args[0] + + +def test_get_api_keys_normalizes_null_allowlist_to_empty_list(): + client = MagicMock() + client.execute.return_value = { + "user": { + "apiKeys": [_api_key_data(allowedIpCidrs=None)], + "apiKeysOtherUsers": [], + } + } + + api_keys = ApiKey.get_api_keys(client) + + assert api_keys[0].allowed_ip_cidrs == [] + + +def test_create_api_key_without_allowlist_sends_null(): + client = MagicMock() + client.execute.return_value = { + "createApiKey": {"id": "key-1", "jwt": "secret"} + } + + with ( + patch.object(ApiKey, "_get_user", return_value="user-1"), + patch.object( + ApiKey, "_get_available_api_key_roles", return_value=["Admin"] + ), + ): + ApiKey.create_api_key( + client, + name="Unrestricted key", + user="person@example.com", + role="Admin", + validity=5, + time_unit=TimeUnit.MINUTE, + ) + + assert client.execute.call_args.args[1]["allowedIpCidrs"] is None diff --git a/libs/labelbox/tests/unit/test_client.py b/libs/labelbox/tests/unit/test_client.py index bbdf0857c..4fc0a0321 100644 --- a/libs/labelbox/tests/unit/test_client.py +++ b/libs/labelbox/tests/unit/test_client.py @@ -1,10 +1,13 @@ -from unittest.mock import Mock +from ipaddress import ip_address, ip_network +from unittest.mock import Mock, patch import pytest from lbox.exceptions import ResourceNotFoundError from labelbox.client import Client +from labelbox.schema.api_key import ApiKey from labelbox.schema.embedding import Embedding +from labelbox.schema.timeunit import TimeUnit # @patch.dict(os.environ, {'LABELBOX_API_KEY': 'bar'}) @@ -121,3 +124,37 @@ def test_embedding_vector_operations_remain_on_adv(): "embedding-id", "vectors.ndjson", callback ) assert embedding.get_imported_vector_count() == 12 + + +def test_create_api_key_forwards_ip_allowlist(): + client = Client(api_key="api_key") + allowed_ip_cidrs = [ + ip_address("203.0.113.10"), + ip_network("198.51.100.0/24"), + ] + with patch.object( + ApiKey, + "create_api_key", + return_value={"id": "key-1", "jwt": "secret"}, + ) as create_api_key: + with pytest.warns(UserWarning, match="currently in alpha"): + result = client.create_api_key( + name="Restricted key", + user="person@example.com", + role="Admin", + validity=5, + time_unit=TimeUnit.MINUTE, + allowed_ip_cidrs=allowed_ip_cidrs, + ) + + create_api_key.assert_called_once_with( + client, + "Restricted key", + "person@example.com", + "Admin", + 5, + TimeUnit.MINUTE, + allowed_ip_cidrs, + ) + + assert result == {"id": "key-1", "jwt": "secret"}