Add pre-defined check is_room_encrypted

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

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