From dd1995ae0b62e1b48dcbffdfe10d044576e41ba8 Mon Sep 17 00:00:00 2001
From: Alero <zhucct34@gmail.com>
Date: Wed, 12 Feb 2025 03:05:51 +0800
Subject: [PATCH] feat: add a way to define custom permission filter.

---
 astrbot/api/event/filter/__init__.py       |  2 ++
 astrbot/core/star/filter/permission.py     | 10 +++++++++-
 astrbot/core/star/register/__init__.py     |  2 ++
 astrbot/core/star/register/star_handler.py | 17 ++++++++++++++++-
 4 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/astrbot/api/event/filter/__init__.py b/astrbot/api/event/filter/__init__.py
index 69e884b89..6dffbb4e6 100644
--- a/astrbot/api/event/filter/__init__.py
+++ b/astrbot/api/event/filter/__init__.py
@@ -5,6 +5,7 @@ from astrbot.core.star.register import (
     register_regex as regex,
     register_platform_adapter_type as platform_adapter_type,
     register_permission_type as permission_type,
+    register_custom_permission_type as custom_permission_type,
     register_on_llm_request as on_llm_request,
     register_on_llm_response as on_llm_response,
     register_llm_tool as llm_tool,
@@ -28,6 +29,7 @@ __all__ = [
     'PlatformAdapterTypeFilter',
     'PlatformAdapterType',
     'PermissionTypeFilter',
+    'register_custom_permission_type',
     'PermissionType',
     'on_llm_request',
     'llm_tool',
diff --git a/astrbot/core/star/filter/permission.py b/astrbot/core/star/filter/permission.py
index 5e961aa5b..23d357338 100644
--- a/astrbot/core/star/filter/permission.py
+++ b/astrbot/core/star/filter/permission.py
@@ -9,7 +9,15 @@ class PermissionType(enum.Flag):
     ADMIN = enum.auto()
     MEMBER = enum.auto()
 
-class PermissionTypeFilter(HandlerFilter):
+class BasePermissionTypeFilter(HandlerFilter):
+    def __init__(self,  raise_error: bool = True):
+        self.raise_error = raise_error
+
+    def filter(self, event: AstrMessageEvent, cfg: AstrBotConfig) -> bool:
+        ''' 一个用于重写的自定义Filter '''
+        raise NotImplementedError
+
+class PermissionTypeFilter(BasePermissionTypeFilter):
     def __init__(self, permission_type: PermissionType, raise_error: bool = True):
         self.permission_type = permission_type
         self.raise_error = raise_error
diff --git a/astrbot/core/star/register/__init__.py b/astrbot/core/star/register/__init__.py
index 422c87b4c..56f40e0e3 100644
--- a/astrbot/core/star/register/__init__.py
+++ b/astrbot/core/star/register/__init__.py
@@ -6,6 +6,7 @@ from .star_handler import (
     register_platform_adapter_type,
     register_regex,
     register_permission_type,
+    register_custom_permission_type,
     register_on_llm_request,
     register_on_llm_response,
     register_llm_tool,
@@ -21,6 +22,7 @@ __all__ = [
     'register_platform_adapter_type',
     'register_regex',
     'register_permission_type',
+    'register_custom_permission_type',
     'register_on_llm_request',
     'register_on_llm_response',
     'register_llm_tool',
diff --git a/astrbot/core/star/register/star_handler.py b/astrbot/core/star/register/star_handler.py
index 0055ce05b..3263eddbe 100644
--- a/astrbot/core/star/register/star_handler.py
+++ b/astrbot/core/star/register/star_handler.py
@@ -6,7 +6,7 @@ from ..filter.command import CommandFilter
 from ..filter.command_group import CommandGroupFilter
 from ..filter.event_message_type import EventMessageTypeFilter, EventMessageType
 from ..filter.platform_adapter_type import PlatformAdapterTypeFilter, PlatformAdapterType
-from ..filter.permission import PermissionTypeFilter, PermissionType
+from ..filter.permission import PermissionTypeFilter, BasePermissionTypeFilter, PermissionType
 from ..filter.regex import RegexFilter
 from typing import Awaitable
 from astrbot.core.provider.func_tool_manager import SUPPORTED_TYPES
@@ -156,6 +156,21 @@ def register_permission_type(permission_type: PermissionType, raise_error: bool
 
     return decorator
 
+def register_custom_permission_type(cunstom_permission_type_filter: BasePermissionTypeFilter, raise_error: bool = True):
+    '''注册一个自定义的 PermissionFilter
+
+    Args:
+        cunstom_permission_type_filter: 一个继承自HandlerFilter
+        raise_error: 如果没有权限，是否抛出错误到消息平台，并且停止事件传播。默认为 True
+    '''
+
+    def decorator(awaitable):
+        handler_md = get_handler_or_create(awaitable, EventType.AdapterMessageEvent)
+        handler_md.event_filters.append(cunstom_permission_type_filter())
+        return awaitable
+
+    return decorator
+
 def register_on_llm_request(**kwargs):
     '''当有 LLM 请求时的事件