1import argparse
2import os
3import sys
4import json
5import time
6import asyncio
7import sqlite3
8import logging
9import mimetypes
10from datetime import datetime, timedelta
11from enum import Enum
12from typing import List, Dict, Optional, Union, Any
13from pathlib import Path
14from urllib.parse import unquote, urlparse
15 
16# Third-party libraries
17import nest_asyncio
18from dotenv import load_dotenv
19from mcp.server.fastmcp import FastMCP, Context
20from mcp.types import ToolAnnotations
21from mcp.shared.exceptions import McpError
22from pythonjsonlogger import jsonlogger
23from telethon import TelegramClient, functions, types, utils
24from telethon.sessions import StringSession
25from telethon.tl.types import (
26    User,
27    Chat,
28    Channel,
29    ChatAdminRights,
30    ChatBannedRights,
31    ChannelParticipantsKicked,
32    ChannelParticipantsAdmins,
33    InputChatPhoto,
34    InputChatUploadedPhoto,
35    InputChatPhotoEmpty,
36    InputPeerUser,
37    InputPeerChat,
38    InputPeerChannel,
39    DialogFilter,
40    DialogFilterChatlist,
41    DialogFilterDefault,
42    TextWithEntities,
43)
44import re
45from functools import wraps
46import telethon.errors.rpcerrorlist
47 
48 
49class ValidationError(Exception):
50    """Custom exception for validation errors."""
51 
52    pass
53 
54 
55def json_serializer(obj):
56    """Helper function to convert non-serializable objects for JSON serialization."""
57    if isinstance(obj, datetime):
58        return obj.isoformat()
59    if isinstance(obj, bytes):
60        return obj.decode("utf-8", errors="replace")
61    # Add other non-serializable types as needed
62    raise TypeError(f"Object of type {type(obj)} is not JSON serializable")
63 
64 
65def get_entity_type(entity: Any) -> str:
66    """Return a normalized, human-readable chat/entity type."""
67    if isinstance(entity, User):
68        return "User"
69    if isinstance(entity, Chat):
70        return "Group (Basic)"
71    if isinstance(entity, Channel):
72        if getattr(entity, "megagroup", False):
73            return "Supergroup"
74        return "Channel" if getattr(entity, "broadcast", False) else "Group"
75    return type(entity).__name__
76 
77 
78def get_entity_filter_type(entity: Any) -> Optional[str]:
79    """Return list_chats-compatible filter type: user/group/channel."""
80    entity_type = get_entity_type(entity)
81    if entity_type == "User":
82        return "user"
83    if entity_type in ("Group (Basic)", "Group", "Supergroup"):
84        return "group"
85    if entity_type == "Channel":
86        return "channel"
87    return None
88 
89 
90load_dotenv()
91 
92TELEGRAM_API_ID = int(os.getenv("TELEGRAM_API_ID"))
93TELEGRAM_API_HASH = os.getenv("TELEGRAM_API_HASH")
94TELEGRAM_SESSION_NAME = os.getenv("TELEGRAM_SESSION_NAME")
95 
96# Check if a string session exists in environment, otherwise use file-based session
97SESSION_STRING = os.getenv("TELEGRAM_SESSION_STRING")
98 
99mcp = FastMCP("telegram")
100 
101if SESSION_STRING:
102    # Use the string session if available
103    client = TelegramClient(StringSession(SESSION_STRING), TELEGRAM_API_ID, TELEGRAM_API_HASH)
104else:
105    # Use file-based session
106    client = TelegramClient(TELEGRAM_SESSION_NAME, TELEGRAM_API_ID, TELEGRAM_API_HASH)
107 
108# Setup robust logging with both file and console output
109logger = logging.getLogger("telegram_mcp")
110logger.setLevel(logging.ERROR)  # Set to ERROR for production, INFO for debugging
111 
112# Create console handler
113console_handler = logging.StreamHandler()
114console_handler.setLevel(logging.ERROR)  # Set to ERROR for production, INFO for debugging
115 
116# Create file handler with absolute path
117script_dir = os.path.dirname(os.path.abspath(__file__))
118log_file_path = os.path.join(script_dir, "mcp_errors.log")
119 
120try:
121    file_handler = logging.FileHandler(log_file_path, mode="a")  # Append mode
122    file_handler.setLevel(logging.ERROR)
123 
124    # Create formatters
125    # Console formatter remains in the old format
126    console_formatter = logging.Formatter("%(asctime)s [%(levelname)s] %(name)s - %(message)s")
127    console_handler.setFormatter(console_formatter)
128 
129    # File formatter is now JSON
130    json_formatter = jsonlogger.JsonFormatter(
131        "%(asctime)s %(name)s %(levelname)s %(message)s",
132        datefmt="%Y-%m-%dT%H:%M:%S%z",
133    )
134    file_handler.setFormatter(json_formatter)
135 
136    # Add handlers to logger
137    logger.addHandler(console_handler)
138    logger.addHandler(file_handler)
139    logger.info(f"Logging initialized to {log_file_path}")
140except Exception as log_error:
141    print(f"WARNING: Error setting up log file: {log_error}", file=sys.stderr)
142    # Fallback to console-only logging
143    logger.addHandler(console_handler)
144    logger.error(f"Failed to set up log file handler: {log_error}")
145 
146 
147# File-path tool security configuration
148SERVER_ALLOWED_ROOTS: list[Path] = []
149DEFAULT_DOWNLOAD_SUBDIR = "downloads"
150DISALLOWED_PATH_PATTERNS = ("*", "?", "[", "]", "{", "}", "~", "\x00")
151EXTENSION_ALLOWLISTS: dict[str, set[str]] = {
152    "send_voice": {".ogg", ".opus"},
153    "send_sticker": {".webp"},
154    "set_profile_photo": {".jpg", ".jpeg", ".png", ".webp"},
155    "edit_chat_photo": {".jpg", ".jpeg", ".png", ".webp"},
156}
157MAX_FILE_BYTES: dict[str, int] = {
158    "send_file": 200 * 1024 * 1024,  # 200 MB
159    "upload_file": 200 * 1024 * 1024,
160    "send_voice": 100 * 1024 * 1024,
161    "send_sticker": 10 * 1024 * 1024,
162    "set_profile_photo": 50 * 1024 * 1024,
163    "edit_chat_photo": 50 * 1024 * 1024,
164}
165ROOTS_UNSUPPORTED_ERROR_CODES = {-32601}
166ROOTS_STATUS_READY = "ready"
167ROOTS_STATUS_NOT_CONFIGURED = "not_configured"
168ROOTS_STATUS_UNSUPPORTED_FALLBACK = "unsupported_fallback"
169ROOTS_STATUS_CLIENT_DENY_ALL = "client_deny_all"
170ROOTS_STATUS_ERROR = "error"
171 
172 
173# Error code prefix mapping for better error tracing
174class ErrorCategory(str, Enum):
175    CHAT = "CHAT"
176    MSG = "MSG"
177    CONTACT = "CONTACT"
178    GROUP = "GROUP"
179    MEDIA = "MEDIA"
180    PROFILE = "PROFILE"
181    AUTH = "AUTH"
182    ADMIN = "ADMIN"
183    FOLDER = "FOLDER"
184 
185 
186def log_and_format_error(
187    function_name: str,
188    error: Exception,
189    prefix: Optional[Union[ErrorCategory, str]] = None,
190    user_message: str = None,
191    **kwargs,
192) -> str:
193    """
194    Centralized error handling function.
195 
196    Logs an error and returns a formatted, user-friendly message.
197 
198    Args:
199        function_name: Name of the function where the error occurred.
200        error: The exception that was raised.
201        prefix: Error code prefix (e.g., ErrorCategory.CHAT, "VALIDATION-001").
202            If None, it will be derived from the function_name.
203        user_message: A custom user-facing message to return. If None, a generic one is created.
204        **kwargs: Additional context parameters to include in the log.
205 
206    Returns:
207        A user-friendly error message with an error code.
208    """
209    # Generate a consistent error code
210    if isinstance(prefix, str) and prefix == "VALIDATION-001":
211        # Special case for validation errors
212        error_code = prefix
213    else:
214        if prefix is None:
215            # Try to derive prefix from function name
216            for category in ErrorCategory:
217                if category.name.lower() in function_name.lower():
218                    prefix = category
219                    break
220 
221        prefix_str = prefix.value if isinstance(prefix, ErrorCategory) else (prefix or "GEN")
222        error_code = f"{prefix_str}-ERR-{abs(hash(function_name)) % 1000:03d}"
223 
224    # Format the additional context parameters
225    context = ", ".join(f"{k}={v}" for k, v in kwargs.items())
226 
227    # Log the full technical error
228    logger.error(f"Error in {function_name} ({context}) - Code: {error_code}", exc_info=True)
229 
230    # Return a user-friendly message
231    if user_message:
232        return user_message
233 
234    return f"An error occurred (code: {error_code}). Check mcp_errors.log for details."
235 
236 
237def validate_id(*param_names_to_validate):
238    """
239    Decorator to validate chat_id and user_id parameters, including lists of IDs.
240    It checks for valid integer ranges, string representations of integers,
241    and username formats.
242    """
243 
244    def decorator(func):
245        @wraps(func)
246        async def wrapper(*args, **kwargs):
247            for param_name in param_names_to_validate:
248                if param_name not in kwargs or kwargs[param_name] is None:
249                    continue
250 
251                param_value = kwargs[param_name]
252 
253                def validate_single_id(value, p_name):
254                    # Handle integer IDs
255                    if isinstance(value, int):
256                        if not (-(2**63) <= value <= 2**63 - 1):
257                            return (
258                                None,
259                                f"Invalid {p_name}: {value}. ID is out of the valid integer range.",
260                            )
261                        return value, None
262 
263                    # Handle string IDs
264                    if isinstance(value, str):
265                        try:
266                            int_value = int(value)
267                            if not (-(2**63) <= int_value <= 2**63 - 1):
268                                return (
269                                    None,
270                                    f"Invalid {p_name}: {value}. ID is out of the valid integer range.",
271                                )
272                            return int_value, None
273                        except ValueError:
274                            if re.match(r"^@?[a-zA-Z0-9_]{5,}$", value):
275                                return value, None
276                            else:
277                                return (
278                                    None,
279                                    f"Invalid {p_name}: '{value}'. Must be a valid integer ID, or a username string.",
280                                )
281 
282                    # Handle other invalid types
283                    return (
284                        None,
285                        f"Invalid {p_name}: {value}. Type must be an integer or a string.",
286                    )
287 
288                if isinstance(param_value, list):
289                    validated_list = []
290                    for item in param_value:
291                        validated_item, error_msg = validate_single_id(item, param_name)
292                        if error_msg:
293                            return log_and_format_error(
294                                func.__name__,
295                                ValidationError(error_msg),
296                                prefix="VALIDATION-001",
297                                user_message=error_msg,
298                                **{param_name: param_value},
299                            )
300                        validated_list.append(validated_item)
301                    kwargs[param_name] = validated_list
302                else:
303                    validated_value, error_msg = validate_single_id(param_value, param_name)
304                    if error_msg:
305                        return log_and_format_error(
306                            func.__name__,
307                            ValidationError(error_msg),
308                            prefix="VALIDATION-001",
309                            user_message=error_msg,
310                            **{param_name: param_value},
311                        )
312                    kwargs[param_name] = validated_value
313 
314            return await func(*args, **kwargs)
315 
316        return wrapper
317 
318    return decorator
319 
320 
321def format_entity(entity) -> Dict[str, Any]:
322    """Helper function to format entity information consistently."""
323    result = {"id": entity.id}
324 
325    if hasattr(entity, "title"):
326        result["name"] = entity.title
327        result["type"] = "group" if isinstance(entity, Chat) else "channel"
328    elif hasattr(entity, "first_name"):
329        name_parts = []
330        if entity.first_name:
331            name_parts.append(entity.first_name)
332        if hasattr(entity, "last_name") and entity.last_name:
333            name_parts.append(entity.last_name)
334        result["name"] = " ".join(name_parts)
335        result["type"] = "user"
336        if hasattr(entity, "username") and entity.username:
337            result["username"] = entity.username
338        if hasattr(entity, "phone") and entity.phone:
339            result["phone"] = entity.phone
340 
341    return result
342 
343 
344async def resolve_entity(identifier: Union[int, str]) -> Any:
345    """Resolve entity with automatic cache warming on miss.
346 
347    StringSession has no persistent entity cache. If get_entity() fails
348    because the cache is cold (ValueError on PeerUser lookup for group IDs),
349    warm the cache via get_dialogs() and retry.
350    """
351    try:
352        return await client.get_entity(identifier)
353    except ValueError:
354        await client.get_dialogs()
355        return await client.get_entity(identifier)
356 
357 
358async def resolve_input_entity(identifier: Union[int, str]) -> Any:
359    """Like resolve_entity() but returns an InputPeer."""
360    try:
361        return await client.get_input_entity(identifier)
362    except ValueError:
363        await client.get_dialogs()
364        return await client.get_input_entity(identifier)
365 
366 
367def format_message(message) -> Dict[str, Any]:
368    """Helper function to format message information consistently."""
369    result = {
370        "id": message.id,
371        "date": message.date.isoformat(),
372        "text": message.message or "",
373    }
374 
375    if message.from_id:
376        result["from_id"] = utils.get_peer_id(message.from_id)
377 
378    if message.media:
379        result["has_media"] = True
380        result["media_type"] = type(message.media).__name__
381 
382    return result
383 
384 
385def get_sender_name(message) -> str:
386    """Helper function to get sender name from a message."""
387    if not message.sender:
388        return "Unknown"
389 
390    # Check for group/channel title first
391    if hasattr(message.sender, "title") and message.sender.title:
392        return message.sender.title
393    elif hasattr(message.sender, "first_name"):
394        # User sender
395        first_name = getattr(message.sender, "first_name", "") or ""
396        last_name = getattr(message.sender, "last_name", "") or ""
397        full_name = f"{first_name} {last_name}".strip()
398        return full_name if full_name else "Unknown"
399    else:
400        return "Unknown"
401 
402 
403def get_engagement_info(message) -> str:
404    """Helper function to get engagement metrics (views, forwards, reactions) from a message."""
405    engagement_parts = []
406    views = getattr(message, "views", None)
407    if views is not None:
408        engagement_parts.append(f"views:{views}")
409    forwards = getattr(message, "forwards", None)
410    if forwards is not None:
411        engagement_parts.append(f"forwards:{forwards}")
412    reactions = getattr(message, "reactions", None)
413    if reactions is not None:
414        results = getattr(reactions, "results", None)
415        total_reactions = sum(getattr(r, "count", 0) or 0 for r in results) if results else 0
416        engagement_parts.append(f"reactions:{total_reactions}")
417    return f" | {', '.join(engagement_parts)}" if engagement_parts else ""
418 
419 
420def _dedupe_paths(paths: List[Path]) -> List[Path]:
421    seen: set[str] = set()
422    result: List[Path] = []
423    for path in paths:
424        key = str(path)
425        if key in seen:
426            continue
427        seen.add(key)
428        result.append(path)
429    return result
430 
431 
432def _contains_forbidden_path_patterns(raw_path: str) -> Optional[str]:
433    value = raw_path.strip()
434    if not value:
435        return "Path must not be empty."
436    if any(token in value for token in DISALLOWED_PATH_PATTERNS):
437        return "Path contains disallowed wildcard/shell patterns."
438    if ".." in Path(value).parts:
439        return "Path traversal is not allowed."
440    return None
441 
442 
443def _coerce_root_uri_to_path(uri: str) -> Path:
444    parsed = urlparse(uri)
445    if parsed.scheme != "file":
446        raise ValueError(f"Unsupported root URI scheme: {parsed.scheme}")
447 
448    decoded_path = unquote(parsed.path or "")
449    if parsed.netloc and parsed.netloc not in ("", "localhost"):
450        decoded_path = f"//{parsed.netloc}{decoded_path}"
451    if os.name == "nt" and decoded_path.startswith("/") and len(decoded_path) > 2:
452        # file:///C:/tmp -> C:/tmp on Windows
453        if decoded_path[2] == ":":
454            decoded_path = decoded_path