From b2d4158637fbda4604497969a66f1c07faeb8377 Mon Sep 17 00:00:00 2001
From: ag5250 <agarodia98@gmail.com>
Date: Fri, 24 Jul 2026 08:11:19 +0000
Subject: [PATCH] Add pre-defined check is_room_encrypted

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 matrix/checks.py     | 28 +++++++++++++++++
 tests/test_checks.py | 73 +++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 100 insertions(+), 1 deletion(-)

diff --git a/matrix/checks.py b/matrix/checks.py
index a3e5f2c..1f64f53 100644
--- a/matrix/checks.py
+++ b/matrix/checks.py
@@ -92,3 +92,31 @@ def wrapper(cmd: "Command") -> "Command":
         return cmd
 
     return wrapper
+
+
+def is_room_encrypted() -> Callable:
+    """
+    Decorator to restrict a command to encrypted rooms.
+
+    ## Example
+
+    ```python
+    @is_room_encrypted()
+    @bot.command("secret")
+    async def secret(ctx: Context) -> None:
+        await ctx.reply("This room is encrypted!")
+
+    @secret.error(CheckError)
+    async def secret_error(ctx: Context, error: CheckError) -> None:
+        await ctx.reply("This command can only be used in an encrypted room.")
+    ```
+    """
+
+    async def _is_room_encrypted(ctx: "Context") -> bool:
+        return ctx.room.encrypted
+
+    def wrapper(cmd: "Command") -> "Command":
+        cmd.check(_is_room_encrypted)
+        return cmd
+
+    return wrapper
diff --git a/tests/test_checks.py b/tests/test_checks.py
index fe70e46..9fb85a3 100644
--- a/tests/test_checks.py
+++ b/tests/test_checks.py
@@ -8,6 +8,7 @@
     MODERATOR_POWER_LEVEL,
     is_admin,
     is_moderator,
+    is_room_encrypted,
 )
 from matrix.command import Command
 from matrix.context import Context
@@ -42,11 +43,14 @@ def make_event(sender: str) -> RoomMessageText:
     )
 
 
-def make_context(bot, client, sender: str, level: int) -> Context:
+def make_context(
+    bot, client, sender: str, level: int, encrypted: bool = False
+) -> Context:
     """Builds a real Context whose room reports `level` as the sender's power level."""
     matrix_room = Mock(spec=MatrixRoom)
     matrix_room.room_id = "!room:example.com"
     matrix_room.power_levels = PowerLevels(users={sender: level})
+    matrix_room.encrypted = encrypted
 
     room = Room(matrix_room, client)
     return Context(bot, room, make_event(sender))
@@ -181,6 +185,73 @@ async def restricted(ctx):
     assert called is True
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("encrypted", [True, False])
+async def test_is_room_encrypted_check__reflects_room_encryption(
+    bot, client, encrypted
+):
+    async def my_command(ctx):
+        pass
+
+    cmd = Command(my_command)
+    is_room_encrypted()(cmd)
+
+    check = cmd.checks[-1]
+    ctx = make_context(bot, client, "@user:example.com", 0, encrypted=encrypted)
+
+    assert await check(ctx) is encrypted
+
+
+def test_is_room_encrypted__returns_the_same_command():
+    async def my_command(ctx):
+        pass
+
+    cmd = Command(my_command)
+    assert is_room_encrypted()(cmd) is cmd
+
+
+@pytest.mark.asyncio
+async def test_is_room_encrypted__raises_check_error_when_not_encrypted(bot, client):
+    called = False
+
+    async def restricted(ctx):
+        nonlocal called
+        called = True
+
+    cmd = Command(restricted)
+    is_room_encrypted()(cmd)
+
+    caught: list[Exception] = []
+
+    @cmd.error(CheckError)
+    async def on_check_error(ctx, error):
+        caught.append(error)
+
+    ctx = make_context(bot, client, "@user:example.com", 0, encrypted=False)
+    await cmd(ctx)
+
+    assert called is False
+    assert len(caught) == 1
+    assert isinstance(caught[0], CheckError)
+
+
+@pytest.mark.asyncio
+async def test_is_room_encrypted__allows_command_when_encrypted(bot, client):
+    called = False
+
+    async def restricted(ctx):
+        nonlocal called
+        called = True
+
+    cmd = Command(restricted)
+    is_room_encrypted()(cmd)
+
+    ctx = make_context(bot, client, "@user:example.com", 0, encrypted=True)
+    await cmd(ctx)
+
+    assert called is True
+
+
 def test_power_level_constants__match_matrix_defaults():
     """The thresholds follow the defaults of `m.room.power_levels`."""
     assert ADMIN_POWER_LEVEL == 100