✨ SubcommandPermExtension

Author: RF-Tar-Railt

src/nonebot_plugin_alconna/builtins/extensions/permission.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+import asyncio
+from typing import Protocol
+
+from arclet.alconna import Arparma, CompSession, SubcommandResult
+from nonebot.internal.adapter import Bot, Event
+
+from nonebot_plugin_alconna import Extension
+
+
+class Checker(Protocol):
+    async def __call__(self, bot: Bot, event: Event, permission: str) -> bool: ...
+
+
+class SubcommandPermExtension(Extension):
+    """
+    用于简易检查调用者是否有命令权限的拓展。
+
+    Example:
+        >>> from nonebot_plugin_alconna.builtins.extensions.permission import SubcommandPermExtension
+        >>>
+        >>> matcher = on_alconna("...", extensions=[SubcommandPermExtension(...)])
+    """
+
+    @property
+    def priority(self) -> int:
+        return 20
+
+    def __init__(self, checker: Checker, include_options: bool = False) -> None:
+        """
+        Args:
+            checker: 权限检查函数，接受 bot、event、permission 三个参数，返回是否有权限的布尔值
+            include_options: 是否需要选项的权限检查
+        """
+        self.checker = checker
+        self.include_options = include_options
+
+    @property
+    def id(self) -> str:
+        return "builtins.extensions.permission:SubcommandPermExtension"
+
+    async def permission_check(self, bot: Bot, event: Event, medium: Arparma | CompSession) -> bool:
+        if isinstance(medium, CompSession):
+            return True
+        base = [f"command.{medium.source.name}"]
+        if self.include_options:
+            base.extend(f"command.{medium.source.name}.$options.{opt}" for opt in medium.options)
+
+        def gen_permissions(subcommands: dict[str, SubcommandResult], prefix: str):
+            for name, result in subcommands.items():
+                current_perm = f"{prefix}.{name}"
+                yield current_perm
+                if self.include_options and result.options:
+                    for opt in result.options:
+                        yield f"{current_perm}.$options.{opt}"
+                if result.subcommands:
+                    yield from gen_permissions(result.subcommands, current_perm)
+
+        base.extend(gen_permissions(medium.subcommands, f"command.{medium.source.name}"))
+        tasks = [self.checker(bot, event, perm) for perm in base]
+        results = await asyncio.gather(*tasks)
+        return all(results)
+
+
+__extension__ = SubcommandPermExtension

src/nonebot_plugin_alconna/extension.py

@@ -11,7 +11,7 @@
 from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, TypeVar, Union, final, overload
 from weakref import finalize
 
-from arclet.alconna import Alconna, Arparma
+from arclet.alconna import Alconna, Arparma, CompSession
 from nonebot import get_plugin_config
 from nonebot.adapters import Bot, Event, Message
 from nonebot.compat import PydanticUndefined
@@ -133,7 +133,7 @@ async def receive_wrapper(self, bot: Bot, event: Event, command: Alconna, receiv
         """接收消息后的钩子函数。"""
         return receive
 
-    async def permission_check(self, bot: Bot, event: Event, command: Alconna) -> bool:
+    async def permission_check(self, bot: Bot, event: Event, medium: Arparma | CompSession) -> bool:
         """命令首次解析并确认头部匹配（即确认选择响应）时对发送者的权限判断"""
         return True
 
@@ -220,10 +220,10 @@ async def receive_wrapper(self, bot: Bot, event: Event, command: Alconna, receiv
                 res = await ext.receive_wrapper(bot, event, command, res)
         return res
 
-    async def permission_check(self, bot: Bot, event: Event, command: Alconna) -> bool:
+    async def permission_check(self, bot: Bot, event: Event, medium: Arparma | CompSession) -> bool:
         for ext in self.context:
             if ext._overrides["permission_check"]:
-                if await ext.permission_check(bot, event, command) is False:
+                if await ext.permission_check(bot, event, medium) is False:
                     return False
                 continue
         return True

src/nonebot_plugin_alconna/rule.py

@@ -241,7 +241,7 @@ async def handle(
         if res:
             interface.exit()
             return res
-        if not await selected.permission_check(bot, event, cmd):
+        if not await selected.permission_check(bot, event, interface):
             return False
 
         res = Arparma(
@@ -354,7 +354,7 @@ async def __call__(
             return False
         if self.skip and may_help_text:
             return False
-        if not await selected.permission_check(bot, event, cmd):
+        if not await selected.permission_check(bot, event, arp):
             return False
         state[ALCONNA_RESULT] = CommandResult(result=arp, output=may_help_text)
         state[ALCONNA_EXEC_RESULT] = cmd.exec_result

tests/test_extension.py

@@ -1,6 +1,6 @@
 from typing import Any
 
-from arclet.alconna import Alconna, Args, Arparma
+from arclet.alconna import Alconna, Args, Arparma, Subcommand
 from nonebot import get_adapter
 from nonebot.internal.adapter import Event
 from nonebot.params import Depends
@@ -35,7 +35,7 @@ def priority(self) -> int:
         def id(self) -> str:
             return "demo"
 
-        async def permission_check(self, bot, event, command):
+        async def permission_check(self, bot, event, medium):
             return await self.inject(check_dep)
 
         def before_catch(self, name: str, annotation: Any, default: Any) -> bool:
@@ -79,3 +79,77 @@ async def h(a: float, b: float, hello: str, world: str, test: str):
         event = fake_group_message_event_v11(message=Message("add 1.3 2.4"), user_id=456)
         ctx.receive_event(bot, event)
         ctx.should_call_send(event, "权限不足！")
+
+
+@pytest.mark.asyncio()
+async def test_extension_permission(app: App):
+    from nonebot.adapters.onebot.v11 import Adapter, Bot, Message
+
+    from nonebot_plugin_alconna import on_alconna
+    from nonebot_plugin_alconna.builtins.extensions.permission import SubcommandPermExtension
+
+    cmd = Alconna(
+        "calc",
+        Subcommand(
+            "add",
+            Args["a", float]["b", float],
+        ),
+        Subcommand(
+            "mul",
+            Args["a", float]["b", float],
+        ),
+        Subcommand(
+            "div",
+            Args["a", float]["b", float],
+        ),
+        Subcommand(
+            "sub",
+            Args["a", float]["b", float],
+        ),
+    )
+
+    async def checker(bot, event, permission):
+        user_id = event.get_user_id()
+        if user_id == "123":
+            return {
+                "command.calc": True,
+                "command.calc.add": True,
+                "command.calc.sub": True,
+                "command.calc.mul": False,
+                "command.calc.div": False,
+            }[permission]
+        if user_id == "456":
+            return {
+                "command.calc": True,
+                "command.calc.add": False,
+                "command.calc.sub": True,
+                "command.calc.mul": False,
+                "command.calc.div": True,
+            }[permission]
+        return True
+
+    mat = on_alconna(cmd, extensions=[SubcommandPermExtension(checker)])
+
+    @mat.handle()
+    async def h(a: float, b: float):
+        await mat.send(f"Result: {a} and {b}")
+
+    async with app.test_matcher(mat) as ctx:  # type: ignore
+        adapter = get_adapter(Adapter)
+        bot = ctx.create_bot(base=Bot, adapter=adapter)
+
+        event = fake_group_message_event_v11(message=Message("calc add 1 2"), user_id=123)
+        ctx.receive_event(bot, event)
+        ctx.should_call_send(event, "Result: 1.0 and 2.0")
+
+        event = fake_group_message_event_v11(message=Message("calc div 1 2"), user_id=123)
+        ctx.receive_event(bot, event)
+        ctx.should_not_pass_rule()
+
+        event = fake_group_message_event_v11(message=Message("calc div 5 3"), user_id=456)
+        ctx.receive_event(bot, event)
+        ctx.should_call_send(event, "Result: 5.0 and 3.0")
+
+        event = fake_group_message_event_v11(message=Message("calc add 1 2"), user_id=456)
+        ctx.receive_event(bot, event)
+        ctx.should_not_pass_rule()