1import os
2from pathlib import Path
3 
4import pytest
5from mcp import types
6from mcp.shared.exceptions import McpError
7from mcp.types import ErrorData
8 
9os.environ.setdefault("TELEGRAM_API_ID", "12345")
10os.environ.setdefault("TELEGRAM_API_HASH", "dummy_hash")
11 
12import main
13 
14 
15class _DummySession:
16    def __init__(self, roots):
17        self._roots = roots
18 
19    async def list_roots(self):
20        return types.ListRootsResult(roots=self._roots)
21 
22 
23class _DummyContext:
24    def __init__(self, roots):
25        self.session = _DummySession(roots)
26 
27 
28class _FailingSession:
29    def __init__(self, error):
30        self._error = error
31 
32    async def list_roots(self):
33        raise self._error
34 
35 
36class _FailingContext:
37    def __init__(self, error):
38        self.session = _FailingSession(error)
39 
40 
41class _MissingRootsSession:
42    pass
43 
44 
45class _MissingRootsContext:
46    def __init__(self):
47        self.session = _MissingRootsSession()
48 
49 
50@pytest.mark.asyncio
51async def test_readable_relative_path_resolves_inside_first_server_root(tmp_path, monkeypatch):
52    root = (tmp_path / "root").resolve()
53    root.mkdir(parents=True)
54    target = root / "document.txt"
55    target.write_text("ok", encoding="utf-8")
56 
57    monkeypatch.setattr(main, "SERVER_ALLOWED_ROOTS", [root])
58 
59    resolved, error = await main._resolve_readable_file_path(
60        raw_path="document.txt",
61        ctx=None,
62        tool_name="send_file",
63    )
64 
65    assert error is None
66    assert resolved == target.resolve()
67 
68 
69@pytest.mark.asyncio
70async def test_readable_path_rejects_traversal(tmp_path, monkeypatch):
71    root = (tmp_path / "root").resolve()
72    root.mkdir(parents=True)
73    monkeypatch.setattr(main, "SERVER_ALLOWED_ROOTS", [root])
74 
75    resolved, error = await main._resolve_readable_file_path(
76        raw_path="../etc/passwd",
77        ctx=None,
78        tool_name="send_file",
79    )
80 
81    assert resolved is None
82    assert error == "Path traversal is not allowed."
83 
84 
85@pytest.mark.asyncio
86async def test_readable_path_rejects_outside_root(tmp_path, monkeypatch):
87    root = (tmp_path / "root").resolve()
88    outside_root = (tmp_path / "outside").resolve()
89    root.mkdir(parents=True)
90    outside_root.mkdir(parents=True)
91 
92    outside_file = outside_root / "outside.txt"
93    outside_file.write_text("no", encoding="utf-8")
94 
95    monkeypatch.setattr(main, "SERVER_ALLOWED_ROOTS", [root])
96 
97    resolved, error = await main._resolve_readable_file_path(
98        raw_path=str(outside_file),
99        ctx=None,
100        tool_name="send_file",
101    )
102 
103    assert resolved is None
104    assert error == "Path is outside allowed roots."
105 
106 
107@pytest.mark.asyncio
108async def test_client_roots_replace_server_allowlist(tmp_path, monkeypatch):
109    server_root = (tmp_path / "server_root").resolve()
110    client_root = (tmp_path / "client_root").resolve()
111    server_root.mkdir(parents=True)
112    client_root.mkdir(parents=True)
113 
114    (server_root / "server.txt").write_text("server", encoding="utf-8")
115    client_file = client_root / "client.txt"
116    client_file.write_text("client", encoding="utf-8")
117 
118    monkeypatch.setattr(main, "SERVER_ALLOWED_ROOTS", [server_root])
119    ctx = _DummyContext([types.Root(uri=client_root.as_uri())])
120 
121    roots = await main._get_effective_allowed_roots(ctx)
122    assert roots == [client_root]
123 
124    resolved, error = await main._resolve_readable_file_path(
125        raw_path="client.txt",
126        ctx=ctx,
127        tool_name="send_file",
128    )
129    assert error is None
130    assert resolved == client_file.resolve()
131 
132 
133@pytest.mark.asyncio
134async def test_empty_client_roots_disable_file_tools(tmp_path, monkeypatch):
135    server_root = (tmp_path / "server_root").resolve()
136    server_root.mkdir(parents=True)
137 
138    monkeypatch.setattr(main, "SERVER_ALLOWED_ROOTS", [server_root])
139    ctx = _DummyContext([])
140 
141    roots = await main._get_effective_allowed_roots(ctx)
142    assert roots == []
143 
144    resolved, error = await main._resolve_readable_file_path(
145        raw_path="server.txt",
146        ctx=ctx,
147        tool_name="send_file",
148    )
149    assert resolved is None
150    assert error is not None
151    assert "empty MCP Roots list" in error
152    assert "deny-all" in error
153 
154 
155@pytest.mark.asyncio
156async def test_mcp_method_not_found_falls_back_to_server_allowlist(tmp_path, monkeypatch):
157    server_root = (tmp_path / "server_root").resolve()
158    server_root.mkdir(parents=True)
159 
160    monkeypatch.setattr(main, "SERVER_ALLOWED_ROOTS", [server_root])
161    ctx = _FailingContext(McpError(ErrorData(code=-32601, message="Method not found")))
162 
163    roots = await main._get_effective_allowed_roots(ctx)
164    assert roots == [server_root]
165 
166 
167@pytest.mark.asyncio
168async def test_missing_list_roots_method_falls_back_to_server_allowlist(tmp_path, monkeypatch):
169    server_root = (tmp_path / "server_root").resolve()
170    server_root.mkdir(parents=True)
171 
172    monkeypatch.setattr(main, "SERVER_ALLOWED_ROOTS", [server_root])
173    ctx = _MissingRootsContext()
174 
175    roots = await main._get_effective_allowed_roots(ctx)
176    assert roots == [server_root]
177 
178 
179@pytest.mark.asyncio
180async def test_unexpected_roots_error_disables_file_path_tools(tmp_path, monkeypatch):
181    server_root = (tmp_path / "server_root").resolve()
182    server_root.mkdir(parents=True)
183    monkeypatch.setattr(main, "SERVER_ALLOWED_ROOTS", [server_root])
184 
185    ctx = _FailingContext(RuntimeError("transport failure"))
186    roots = await main._get_effective_allowed_roots(ctx)
187    assert roots == []
188 
189    resolved, error = await main._resolve_readable_file_path(
190        raw_path="anything.txt",
191        ctx=ctx,
192        tool_name="send_file",
193    )
194    assert resolved is None
195    assert error is not None
196    assert "disabled" in error
197 
198 
199@pytest.mark.asyncio
200async def test_writable_default_path_uses_downloads_subdir(tmp_path, monkeypatch):
201    root = (tmp_path / "root").resolve()
202    root.mkdir(parents=True)
203    monkeypatch.setattr(main, "SERVER_ALLOWED_ROOTS", [root])
204 
205    resolved, error = await main._resolve_writable_file_path(
206        raw_path=None,
207        default_filename="example.bin",
208        ctx=None,
209        tool_name="download_media",
210    )
211 
212    assert error is None
213    assert resolved == (root / "downloads" / "example.bin").resolve()
214    assert resolved.parent.exists()
215 
216 
217@pytest.mark.asyncio
218async def test_extension_allowlist_is_enforced_for_sticker(tmp_path, monkeypatch):
219    root = (tmp_path / "root").resolve()
220    root.mkdir(parents=True)
221    file_path = root / "sticker.txt"
222    file_path.write_text("bad", encoding="utf-8")
223 
224    monkeypatch.setattr(main, "SERVER_ALLOWED_ROOTS", [root])
225 
226    resolved, error = await main._resolve_readable_file_path(
227        raw_path=str(file_path),
228        ctx=None,
229        tool_name="send_sticker",
230    )
231 
232    assert resolved is None
233    assert error is not None
234    assert "extension is not allowed" in error
235 
236 
237@pytest.mark.asyncio
238async def test_file_tools_disabled_without_any_roots(monkeypatch):
239    monkeypatch.setattr(main, "SERVER_ALLOWED_ROOTS", [])
240 
241    resolved, error = await main._resolve_readable_file_path(
242        raw_path="anything.txt",
243        ctx=None,
244        tool_name="send_file",
245    )
246 
247    assert resolved is None
248    assert error is not None
249    assert "disabled" in error
250