From 37fba2296f48a286b80ed4250e97cfa9138b53a5 Mon Sep 17 00:00:00 2001 From: Soulter <905617992@qq.com> Date: Thu, 18 Jun 2026 23:22:44 +0800 Subject: [PATCH] fix: enforce future task owner checks --- astrbot/core/tools/cron_tools.py | 28 +++++++-- tests/unit/test_cron_tools.py | 98 ++++++++++++++++++++++++++++++++ 2 files changed, 121 insertions(+), 5 deletions(-) diff --git a/astrbot/core/tools/cron_tools.py b/astrbot/core/tools/cron_tools.py index bbfa5729d..ee6ff27e2 100644 --- a/astrbot/core/tools/cron_tools.py +++ b/astrbot/core/tools/cron_tools.py @@ -23,6 +23,23 @@ def _extract_job_session(job: Any) -> str | None: return str(session) if session is not None else None +def _extract_job_sender(job: Any) -> str | None: + payload = getattr(job, "payload", None) + if not isinstance(payload, dict): + return None + sender_id = payload.get("sender_id") + return str(sender_id) if sender_id is not None else None + + +def _job_belongs_to_current_sender( + job: Any, current_umo: str, current_sender_id: str +) -> bool: + return ( + _extract_job_session(job) == current_umo + and _extract_job_sender(job) == current_sender_id + ) + + def _parse_run_at(run_at: Any) -> datetime | None: if run_at in (None, ""): return None @@ -133,6 +150,7 @@ class FutureTaskTool(FunctionTool[AstrAgentContext]): return f"Scheduled future task {job.job_id} ({job.name}) {suffix}." current_umo = context.context.event.unified_msg_origin + current_sender_id = str(context.context.event.get_sender_id()) if action == "edit": job_id = kwargs.get("job_id") if not job_id: @@ -146,8 +164,8 @@ class FutureTaskTool(FunctionTool[AstrAgentContext]): job = await cron_mgr.db.get_cron_job(str(job_id)) if not job: return f"error: cron job {job_id} not found." - if _extract_job_session(job) != current_umo: - return "error: you can only edit future tasks in the current umo." + if not _job_belongs_to_current_sender(job, current_umo, current_sender_id): + return "error: you can only edit your own future tasks." payload = dict(job.payload) if isinstance(job.payload, dict) else {} @@ -214,8 +232,8 @@ class FutureTaskTool(FunctionTool[AstrAgentContext]): job = await cron_mgr.db.get_cron_job(str(job_id)) if not job: return f"error: cron job {job_id} not found." - if _extract_job_session(job) != current_umo: - return "error: you can only delete future tasks in the current umo." + if not _job_belongs_to_current_sender(job, current_umo, current_sender_id): + return "error: you can only delete your own future tasks." await cron_mgr.delete_job(str(job_id)) return f"Deleted cron job {job_id}." @@ -223,7 +241,7 @@ class FutureTaskTool(FunctionTool[AstrAgentContext]): jobs = [ job for job in await cron_mgr.list_jobs() - if _extract_job_session(job) == current_umo + if _job_belongs_to_current_sender(job, current_umo, current_sender_id) ] if not jobs: return "No cron jobs found." diff --git a/tests/unit/test_cron_tools.py b/tests/unit/test_cron_tools.py index a0e520d46..2cc254fbe 100644 --- a/tests/unit/test_cron_tools.py +++ b/tests/unit/test_cron_tools.py @@ -8,6 +8,36 @@ import pytest from astrbot.core.tools.cron_tools import FutureTaskTool +def _context(cron_mgr, *, umo: str = "test:group:shared", sender_id: str = "user-1"): + return SimpleNamespace( + context=SimpleNamespace( + context=SimpleNamespace(cron_manager=cron_mgr), + event=SimpleNamespace( + unified_msg_origin=umo, + get_sender_id=lambda: sender_id, + ), + ) + ) + + +def _job(job_id: str, *, umo: str = "test:group:shared", sender_id: str = "user-1"): + return SimpleNamespace( + job_id=job_id, + name=f"name-{job_id}", + job_type="active_agent", + run_once=False, + cron_expression="0 8 * * *", + enabled=True, + next_run_time=None, + payload={ + "session": umo, + "sender_id": sender_id, + "note": f"note-{job_id}", + "origin": "tool", + }, + ) + + def test_future_task_schema_has_action_and_create_cron_guidance(): """The merged tool should expose action routing and unambiguous cron guidance.""" tool = FutureTaskTool() @@ -124,3 +154,71 @@ async def test_future_task_edit_updates_existing_job(): }, ) assert result == "Updated future task job-1 (new name)." + + +@pytest.mark.asyncio +async def test_future_task_edit_rejects_same_umo_different_sender(): + """Same-session users should not edit another sender's task.""" + tool = FutureTaskTool() + existing_job = _job("job-1", sender_id="admin-user") + cron_mgr = SimpleNamespace( + db=SimpleNamespace(get_cron_job=AsyncMock(return_value=existing_job)), + update_job=AsyncMock(), + ) + + result = await tool.call( + _context(cron_mgr, sender_id="attacker-user"), + action="edit", + job_id="job-1", + note="attacker note", + ) + + assert result == "error: you can only edit your own future tasks." + cron_mgr.update_job.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_future_task_delete_rejects_same_umo_different_sender(): + """Same-session users should not delete another sender's task.""" + tool = FutureTaskTool() + existing_job = _job("job-1", sender_id="admin-user") + cron_mgr = SimpleNamespace( + db=SimpleNamespace(get_cron_job=AsyncMock(return_value=existing_job)), + delete_job=AsyncMock(), + ) + + result = await tool.call( + _context(cron_mgr, sender_id="attacker-user"), + action="delete", + job_id="job-1", + ) + + assert result == "error: you can only delete your own future tasks." + cron_mgr.delete_job.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_future_task_list_filters_by_umo_and_sender(): + """List mode should show only tasks owned by the current sender.""" + tool = FutureTaskTool() + own_job = _job("own-job", sender_id="user-1") + same_umo_other_sender = _job("other-sender-job", sender_id="user-2") + different_umo_same_sender = _job( + "other-umo-job", + umo="test:group:other", + sender_id="user-1", + ) + cron_mgr = SimpleNamespace( + list_jobs=AsyncMock( + return_value=[own_job, same_umo_other_sender, different_umo_same_sender] + ) + ) + + result = await tool.call( + _context(cron_mgr, sender_id="user-1"), + action="list", + ) + + assert "own-job" in result + assert "other-sender-job" not in result + assert "other-umo-job" not in result