From 3919f1219a52b312bdc13687d9a0ea04da1795ae Mon Sep 17 00:00:00 2001
From: Mix <32300164+mnixry@users.noreply.github.com>
Date: Mon, 1 Jan 2024 15:47:16 +0800
Subject: [PATCH 1/7] :lock: Prevent accessing private attribute in
 MessageTemplate

to prevent potential information leak caused by format string injection
---
 nonebot/internal/adapter/template.py | 25 ++++++++++++++++++++++++-
 tests/test_adapters/test_template.py | 23 ++++++++++++++++++-----
 2 files changed, 42 insertions(+), 6 deletions(-)

diff --git a/nonebot/internal/adapter/template.py b/nonebot/internal/adapter/template.py
index f0edc32e8893..174d8fa63616 100644
--- a/nonebot/internal/adapter/template.py
+++ b/nonebot/internal/adapter/template.py
@@ -23,6 +23,13 @@
 if TYPE_CHECKING:
     from .message import Message, MessageSegment
 
+    def formatter_field_name_split(
+        field_name: str,
+    ) -> Tuple[str, List[Tuple[bool, str]]]:
+        ...
+else:
+    from _string import formatter_field_name_split
+
 TM = TypeVar("TM", bound="Message")
 TF = TypeVar("TF", str, "Message")
 
@@ -51,11 +58,14 @@ def __init__(
         ...
 
     def __init__(
-        self, template: Union[str, TM], factory: Union[Type[str], Type[TM]] = str
+        self,
+        template: Union[str, TM],
+        factory: Union[Type[str], Type[TM]] = str,
     ) -> None:
         self.template: TF = template  # type: ignore
         self.factory: Type[TF] = factory  # type: ignore
         self.format_specs: Dict[str, FormatSpecFunc] = {}
+        self.private_getattr = False
 
     def __repr__(self) -> str:
         return f"MessageTemplate({self.template!r}, factory={self.factory!r})"
@@ -167,6 +177,19 @@ def _vformat(
 
         return functools.reduce(self._add, results), auto_arg_index
 
+    def get_field(
+        self, field_name: str, args: Sequence[Any], kwargs: Mapping[str, Any]
+    ) -> Any:
+        first, rest = formatter_field_name_split(field_name)
+        obj = self.get_value(first, args, kwargs)
+
+        for is_attr, value in rest:
+            if not self.private_getattr and is_attr and value.startswith("_"):
+                raise ValueError("attribute name must not start with underscore")
+            obj = getattr(obj, value) if is_attr else obj[value]
+
+        return obj, first
+
     def format_field(self, value: Any, format_spec: str) -> Any:
         formatter: Optional[FormatSpecFunc] = self.format_specs.get(format_spec)
         if formatter is None and not issubclass(self.factory, str):
diff --git a/tests/test_adapters/test_template.py b/tests/test_adapters/test_template.py
index 710c556a6600..f81999180916 100644
--- a/tests/test_adapters/test_template.py
+++ b/tests/test_adapters/test_template.py
@@ -1,5 +1,6 @@
 from nonebot.adapters import MessageTemplate
 from utils import FakeMessage, FakeMessageSegment, escape_text
+import pytest
 
 
 def test_template_basis():
@@ -15,12 +16,8 @@ def test_template_message():
     def custom(input: str) -> str:
         return f"{input}-custom!"
 
-    try:
+    with pytest.raises(ValueError, match="already exists"):
         template.add_format_spec(custom)
-    except ValueError:
-        pass
-    else:
-        raise AssertionError("Should raise ValueError")
 
     format_args = {
         "a": "custom",
@@ -57,3 +54,19 @@ def test_message_injection():
     message = template.format(name="[fake:image]")
 
     assert message.extract_plain_text() == escape_text("[fake:image]Is Bad")
+
+
+def test_malformed_template():
+    positive_template = FakeMessage.template("{a}{b}")
+    message = positive_template.format(a="a", b="b")
+    assert message.extract_plain_text() == "ab"
+
+    malformed_template = FakeMessage.template(
+        "{a.__init__}{b[__builtins__][__import__]}"
+    )
+
+    with pytest.raises(ValueError, match="must not start with"):
+        message = malformed_template.format(a="a", b=globals())
+
+    malformed_template.private_getattr = True
+    message = malformed_template.format(a="a", b=globals())

From 59cce769dfa927a483ba2a32c8f3abbf8539a1ba Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 1 Jan 2024 07:59:17 +0000
Subject: [PATCH 2/7] :rotating_light: auto fix by pre-commit hooks

---
 nonebot/internal/adapter/template.py | 1 +
 tests/test_adapters/test_template.py | 3 ++-
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/nonebot/internal/adapter/template.py b/nonebot/internal/adapter/template.py
index 174d8fa63616..dee4ce9c606e 100644
--- a/nonebot/internal/adapter/template.py
+++ b/nonebot/internal/adapter/template.py
@@ -27,6 +27,7 @@ def formatter_field_name_split(
         field_name: str,
     ) -> Tuple[str, List[Tuple[bool, str]]]:
         ...
+
 else:
     from _string import formatter_field_name_split
 
diff --git a/tests/test_adapters/test_template.py b/tests/test_adapters/test_template.py
index f81999180916..1ca24f799cf0 100644
--- a/tests/test_adapters/test_template.py
+++ b/tests/test_adapters/test_template.py
@@ -1,6 +1,7 @@
+import pytest
+
 from nonebot.adapters import MessageTemplate
 from utils import FakeMessage, FakeMessageSegment, escape_text
-import pytest
 
 
 def test_template_basis():

From 9c91c410c9c7ed79932fe199c914f384a8f32d7f Mon Sep 17 00:00:00 2001
From: Mix <32300164+mnixry@users.noreply.github.com>
Date: Mon, 1 Jan 2024 16:05:51 +0800
Subject: [PATCH 3/7] :bug: Fix attribute name validation in MessageTemplate
 class

---
 nonebot/internal/adapter/template.py |  2 +-
 tests/test_adapters/test_template.py | 11 ++++++-----
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/nonebot/internal/adapter/template.py b/nonebot/internal/adapter/template.py
index dee4ce9c606e..a6af5f1f5b85 100644
--- a/nonebot/internal/adapter/template.py
+++ b/nonebot/internal/adapter/template.py
@@ -185,7 +185,7 @@ def get_field(
         obj = self.get_value(first, args, kwargs)
 
         for is_attr, value in rest:
-            if not self.private_getattr and is_attr and value.startswith("_"):
+            if not self.private_getattr and value.startswith("_"):
                 raise ValueError("attribute name must not start with underscore")
             obj = getattr(obj, value) if is_attr else obj[value]
 
diff --git a/tests/test_adapters/test_template.py b/tests/test_adapters/test_template.py
index 1ca24f799cf0..9622306b2ba2 100644
--- a/tests/test_adapters/test_template.py
+++ b/tests/test_adapters/test_template.py
@@ -62,12 +62,13 @@ def test_malformed_template():
     message = positive_template.format(a="a", b="b")
     assert message.extract_plain_text() == "ab"
 
-    malformed_template = FakeMessage.template(
-        "{a.__init__}{b[__builtins__][__import__]}"
-    )
+    malformed_template = FakeMessage.template("{a.__init__}")
+    with pytest.raises(ValueError, match="must not start with"):
+        message = malformed_template.format(a="a")
 
+    malformed_template = FakeMessage.template("{a[__builtins__]}")
     with pytest.raises(ValueError, match="must not start with"):
-        message = malformed_template.format(a="a", b=globals())
+        message = malformed_template.format(a=globals())
 
     malformed_template.private_getattr = True
-    message = malformed_template.format(a="a", b=globals())
+    message = malformed_template.format(a=globals())

From 603eedc0e7c5667f4b132bb699d3873b981b2641 Mon Sep 17 00:00:00 2001
From: Mix <32300164+mnixry@users.noreply.github.com>
Date: Mon, 1 Jan 2024 16:32:15 +0800
Subject: [PATCH 4/7] :art: Add private_getattr parameter to MessageTemplate
 constructor

---
 nonebot/internal/adapter/template.py | 14 +++++++++++---
 tests/test_adapters/test_template.py |  6 ++++--
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/nonebot/internal/adapter/template.py b/nonebot/internal/adapter/template.py
index a6af5f1f5b85..a6117ee54e34 100644
--- a/nonebot/internal/adapter/template.py
+++ b/nonebot/internal/adapter/template.py
@@ -44,17 +44,24 @@ class MessageTemplate(Formatter, Generic[TF]):
     参数:
         template: 模板
         factory: 消息类型工厂，默认为 `str`
+        private_getattr: 是否允许在模板中访问私有属性，默认为 `False`
     """
 
     @overload
     def __init__(
-        self: "MessageTemplate[str]", template: str, factory: Type[str] = str
+        self: "MessageTemplate[str]",
+        template: str,
+        factory: Type[str] = str,
+        private_getattr: bool = False,
     ) -> None:
         ...
 
     @overload
     def __init__(
-        self: "MessageTemplate[TM]", template: Union[str, TM], factory: Type[TM]
+        self: "MessageTemplate[TM]",
+        template: Union[str, TM],
+        factory: Type[TM],
+        private_getattr: bool = False,
     ) -> None:
         ...
 
@@ -62,11 +69,12 @@ def __init__(
         self,
         template: Union[str, TM],
         factory: Union[Type[str], Type[TM]] = str,
+        private_getattr: bool = False,
     ) -> None:
         self.template: TF = template  # type: ignore
         self.factory: Type[TF] = factory  # type: ignore
         self.format_specs: Dict[str, FormatSpecFunc] = {}
-        self.private_getattr = False
+        self.private_getattr = private_getattr
 
     def __repr__(self) -> str:
         return f"MessageTemplate({self.template!r}, factory={self.factory!r})"
diff --git a/tests/test_adapters/test_template.py b/tests/test_adapters/test_template.py
index 9622306b2ba2..4318d03ebc8a 100644
--- a/tests/test_adapters/test_template.py
+++ b/tests/test_adapters/test_template.py
@@ -70,5 +70,7 @@ def test_malformed_template():
     with pytest.raises(ValueError, match="must not start with"):
         message = malformed_template.format(a=globals())
 
-    malformed_template.private_getattr = True
-    message = malformed_template.format(a=globals())
+    malformed_template = MessageTemplate(
+        "{a[__builtins__][__import__]}{b.__init__}", private_getattr=True
+    )
+    message = malformed_template.format(a=globals(), b="b")

From b7058365fb28efe71fe8ed17bc1add223dc1a6a6 Mon Sep 17 00:00:00 2001
From: Mix <32300164+mnixry@users.noreply.github.com>
Date: Mon, 1 Jan 2024 17:07:58 +0800
Subject: [PATCH 5/7] :art: Add formatter_field_name_split import

---
 nonebot/internal/adapter/template.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/nonebot/internal/adapter/template.py b/nonebot/internal/adapter/template.py
index a6117ee54e34..09c1da625192 100644
--- a/nonebot/internal/adapter/template.py
+++ b/nonebot/internal/adapter/template.py
@@ -1,5 +1,7 @@
 import functools
 from string import Formatter
+from _string import formatter_field_name_split  # type: ignore
+
 from typing_extensions import TypeAlias
 from typing import (
     TYPE_CHECKING,
@@ -23,13 +25,11 @@
 if TYPE_CHECKING:
     from .message import Message, MessageSegment
 
-    def formatter_field_name_split(
+    def formatter_field_name_split(  # noqa: F811
         field_name: str,
     ) -> Tuple[str, List[Tuple[bool, str]]]:
         ...
 
-else:
-    from _string import formatter_field_name_split
 
 TM = TypeVar("TM", bound="Message")
 TF = TypeVar("TF", str, "Message")
@@ -188,7 +188,7 @@ def _vformat(
 
     def get_field(
         self, field_name: str, args: Sequence[Any], kwargs: Mapping[str, Any]
-    ) -> Any:
+    ) -> Tuple[Any, Union[int, str]]:
         first, rest = formatter_field_name_split(field_name)
         obj = self.get_value(first, args, kwargs)
 

From fd165750d44d3b581e8d7b8046ad0e1824b1f793 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 1 Jan 2024 09:08:26 +0000
Subject: [PATCH 6/7] :rotating_light: auto fix by pre-commit hooks

---
 nonebot/internal/adapter/template.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/nonebot/internal/adapter/template.py b/nonebot/internal/adapter/template.py
index 09c1da625192..f801fd458ede 100644
--- a/nonebot/internal/adapter/template.py
+++ b/nonebot/internal/adapter/template.py
@@ -1,7 +1,5 @@
 import functools
 from string import Formatter
-from _string import formatter_field_name_split  # type: ignore
-
 from typing_extensions import TypeAlias
 from typing import (
     TYPE_CHECKING,
@@ -22,6 +20,8 @@
     overload,
 )
 
+from _string import formatter_field_name_split  # type: ignore
+
 if TYPE_CHECKING:
     from .message import Message, MessageSegment
 

From ee2dd1ce88c38f65bf1535f593730ae21f4a6575 Mon Sep 17 00:00:00 2001
From: Mix <32300164+mnixry@users.noreply.github.com>
Date: Mon, 1 Jan 2024 17:28:20 +0800
Subject: [PATCH 7/7] :speech_balloon: Update error message accessing private
 field in template

---
 nonebot/internal/adapter/template.py | 2 +-
 tests/test_adapters/test_template.py | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/nonebot/internal/adapter/template.py b/nonebot/internal/adapter/template.py
index f801fd458ede..b720c25796e1 100644
--- a/nonebot/internal/adapter/template.py
+++ b/nonebot/internal/adapter/template.py
@@ -194,7 +194,7 @@ def get_field(
 
         for is_attr, value in rest:
             if not self.private_getattr and value.startswith("_"):
-                raise ValueError("attribute name must not start with underscore")
+                raise ValueError("Cannot access private attribute")
             obj = getattr(obj, value) if is_attr else obj[value]
 
         return obj, first
diff --git a/tests/test_adapters/test_template.py b/tests/test_adapters/test_template.py
index 4318d03ebc8a..3ca840ea0a7b 100644
--- a/tests/test_adapters/test_template.py
+++ b/tests/test_adapters/test_template.py
@@ -63,11 +63,11 @@ def test_malformed_template():
     assert message.extract_plain_text() == "ab"
 
     malformed_template = FakeMessage.template("{a.__init__}")
-    with pytest.raises(ValueError, match="must not start with"):
+    with pytest.raises(ValueError, match="private attribute"):
         message = malformed_template.format(a="a")
 
     malformed_template = FakeMessage.template("{a[__builtins__]}")
-    with pytest.raises(ValueError, match="must not start with"):
+    with pytest.raises(ValueError, match="private attribute"):
         message = malformed_template.format(a=globals())
 
     malformed_template = MessageTemplate(