NAS-129754 / 24.10 / Update api_key.py to use the new API style (#1…

…3957)

* api_key.do_create new api decorator

* give update and delete new api decorator

* type hints for easier integration

* Finalize models, type hints

* `privilege.py` still depends on `allowlist_item` schema

* truthful type hints

* remove unused imports

* remove arbitrary class instance ability from `ApiKeyEntry` model

* remove use of `NonEmptyString`

src/middlewared/middlewared/api/v25_04_0/__init__.py

@@ -1,3 +1,4 @@
+from .api_key import *  # noqa
 from .cloud_sync import *  # noqa
 from .common import *  # noqa
 from .user import *  # noqa

src/middlewared/middlewared/api/v25_04_0/api_key.py

@@ -0,0 +1,60 @@
+from datetime import datetime
+from typing import Literal, TypeAlias
+from typing_extensions import Annotated
+
+from pydantic import ConfigDict, StringConstraints
+
+from middlewared.api.base import BaseModel, Excluded, excluded_field, NonEmptyString, Private
+
+
+HttpVerb: TypeAlias = Literal["GET", "POST", "PUT", "DELETE", "CALL", "SUBSCRIBE", "*"]
+
+
+class AllowListItem(BaseModel):
+    method: HttpVerb
+    resource: NonEmptyString
+
+
+class ApiKeyEntry(BaseModel):
+    """Represents a record in the account.api_key table."""
+
+    id: int
+    name: Annotated[NonEmptyString, StringConstraints(max_length=200)]
+    key: Private[str]
+    created_at: datetime
+    allowlist: list[AllowListItem]
+
+
+class ApiKeyCreate(ApiKeyEntry):
+    id: Excluded = excluded_field()
+    key: Excluded = excluded_field()
+    created_at: Excluded = excluded_field()
+
+
+class ApiKeyCreateArgs(BaseModel):
+    api_key_create: ApiKeyCreate
+
+
+class ApiKeyCreateResult(BaseModel):
+    result: ApiKeyEntry
+
+
+class ApiKeyUpdate(ApiKeyCreate):
+    reset: bool
+
+
+class ApiKeyUpdateArgs(BaseModel):
+    id: int
+    api_key_update: ApiKeyUpdate
+
+
+class ApiKeyUpdateResult(BaseModel):
+    result: ApiKeyEntry
+
+
+class ApiKeyDeleteArgs(BaseModel):
+    id: int
+
+
+class ApiKeyDeleteResult(BaseModel):
+    result: Literal[True]

src/middlewared/middlewared/plugins/account_/privilege.py

@@ -51,7 +51,11 @@ class Config:
         Str("name", required=True, empty=False),
         List("local_groups", items=[Ref("group_entry")]),
         List("ds_groups", items=[Ref("group_entry")]),
-        List("allowlist", items=[Ref("allowlist_item")]),
+        List("allowlist", items=[Dict(
+            "allowlist_item",
+            Str("method", required=True, enum=["GET", "POST", "PUT", "DELETE", "CALL", "SUBSCRIBE", "*"]),
+            Str("resource", required=True),
+        )]),
         List("roles", items=[Str("role")]),
         Bool("web_shell", required=True),
     )

src/middlewared/middlewared/plugins/api_key.py

@@ -1,13 +1,21 @@
 from datetime import datetime
 import random
 import string
+from typing import Literal, TYPE_CHECKING
 
 from passlib.hash import pbkdf2_sha256
 
-from middlewared.schema import accepts, Bool, Dict, Int, List, Str, Patch
+from middlewared.api import api_method
+from middlewared.api.current import (
+    ApiKeyCreateArgs, ApiKeyCreateResult, ApiKeyUpdateArgs,
+    ApiKeyUpdateResult, ApiKeyDeleteArgs, ApiKeyDeleteResult,
+    HttpVerb
+)
 from middlewared.service import CRUDService, private, ValidationErrors
 import middlewared.sqlalchemy as sa
 from middlewared.utils.allowlist import Allowlist
+if TYPE_CHECKING:
+    from middlewared.main import Middleware
 
 
 class APIKeyModel(sa.Model):
@@ -21,45 +29,26 @@ class APIKeyModel(sa.Model):
 
 
 class ApiKey:
-    def __init__(self, api_key):
+    def __init__(self, api_key: dict):
         self.api_key = api_key
         self.allowlist = Allowlist(self.api_key["allowlist"])
 
-    def authorize(self, method, resource):
+    def authorize(self, method: HttpVerb, resource: str) -> bool:
         return self.allowlist.authorize(method, resource)
 
 
 class ApiKeyService(CRUDService):
 
-    keys = {}
+    keys: dict[int, dict] = {}
 
     class Config:
         namespace = "api_key"
         datastore = "account.api_key"
         datastore_extend = "api_key.item_extend"
         cli_namespace = "auth.api_key"
 
-    @private
-    async def item_extend(self, item):
-        item.pop("key")
-        return item
-
-    @accepts(
-        Dict(
-            "api_key_create",
-            Str("name", required=True, empty=False),
-            List("allowlist", items=[
-                Dict(
-                    "allowlist_item",
-                    Str("method", required=True, enum=["GET", "POST", "PUT", "DELETE", "CALL", "SUBSCRIBE", "*"]),
-                    Str("resource", required=True),
-                    register=True,
-                ),
-            ]),
-            register=True,
-        )
-    )
-    async def do_create(self, data):
+    @api_method(ApiKeyCreateArgs, ApiKeyCreateResult)
+    async def do_create(self, data: dict) -> dict:
         """
         Creates API Key.
 
@@ -82,16 +71,8 @@ async def do_create(self, data):
 
         return self._serve(data, key)
 
-    @accepts(
-        Int("id", required=True),
-        Patch(
-            "api_key_create",
-            "api_key_update",
-            ("add", Bool("reset")),
-            ("attr", {"update": True}),
-        )
-    )
-    async def do_update(self, id_, data):
+    @api_method(ApiKeyUpdateArgs, ApiKeyUpdateResult)
+    async def do_update(self, id_: int, data: dict) -> dict:
         """
         Update API Key `id`.
 
@@ -122,10 +103,8 @@ async def do_update(self, id_, data):
 
         return self._serve(await self.get_instance(id_), key)
 
-    @accepts(
-        Int("id")
-    )
-    async def do_delete(self, id_):
+    @api_method(ApiKeyDeleteArgs, ApiKeyDeleteResult)
+    async def do_delete(self, id_: int) -> Literal[True]:
         """
         Delete API Key `id`.
         """
@@ -147,7 +126,7 @@ async def load_keys(self):
         }
 
     @private
-    async def load_key(self, id_):
+    async def load_key(self, id_: int):
         self.keys[id_] = await self.middleware.call(
             "datastore.query",
             "account.api_key",
@@ -156,7 +135,7 @@ async def load_key(self, id_):
         )
 
     @private
-    async def authenticate(self, key):
+    async def authenticate(self, key: str) -> ApiKey | None:
         try:
             key_id, key = key.split("-", 1)
             key_id = int(key_id)
@@ -173,7 +152,7 @@ async def authenticate(self, key):
 
         return ApiKey(db_key)
 
-    async def _validate(self, schema_name, data, id_=None):
+    async def _validate(self, schema_name: str, data: dict, id_: int=None):
         verrors = ValidationErrors()
 
         await self._ensure_unique(verrors, schema_name, "name", data["name"], id_)
@@ -183,12 +162,12 @@ async def _validate(self, schema_name, data, id_=None):
     def _generate(self):
         return "".join([random.SystemRandom().choice(string.ascii_letters + string.digits) for _ in range(64)])
 
-    def _serve(self, data, key):
+    def _serve(self, data: dict, key: str | None) -> dict:
         if key is None:
             return data
 
         return dict(data, key=f"{data['id']}-{key}")
 
 
-async def setup(middleware):
+async def setup(middleware: 'Middleware'):
     await middleware.call("api_key.load_keys")

src/middlewared/middlewared/utils/allowlist.py

@@ -1,14 +1,16 @@
 import fnmatch
 import re
 
+from middlewared.api.current import HttpVerb
+
 ALLOW_LIST_FULL_ADMIN = {'method': '*', 'resource': '*'}
 
 
 class Allowlist:
-    def __init__(self, allowlist):
-        self.exact = {}
+    def __init__(self, allowlist: list[dict]):
+        self.exact: dict[HttpVerb, set[str]] = {}
         self.full_admin = ALLOW_LIST_FULL_ADMIN in allowlist
-        self.patterns = {}
+        self.patterns: dict[HttpVerb, list[re.Pattern]] = {}
         for entry in allowlist:
             method = entry["method"]
             resource = entry["resource"]
@@ -19,10 +21,10 @@ def __init__(self, allowlist):
                 self.exact.setdefault(method, set())
                 self.exact[method].add(resource)
 
-    def authorize(self, method, resource):
+    def authorize(self, method: HttpVerb, resource: str):
         return self._authorize_internal("*", resource) or self._authorize_internal(method, resource)
 
-    def _authorize_internal(self, method, resource):
+    def _authorize_internal(self, method: HttpVerb, resource: str):
         if (exact := self.exact.get(method)) and resource in exact:
             return True