[ALNR-7032] Add API key IP allowlist support

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]:

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: