Add download file service to Telegram bot integration

homeassistant/components/telegram_bot/__init__.py

@@ -46,9 +46,12 @@
     ATTR_CAPTION,
     ATTR_CHAT_ACTION,
     ATTR_CHAT_ID,
+    ATTR_DIRECTORY_PATH,
     ATTR_DISABLE_NOTIF,
     ATTR_DISABLE_WEB_PREV,
     ATTR_FILE,
+    ATTR_FILE_ID,
+    ATTR_FILE_NAME,
     ATTR_IS_ANONYMOUS,
     ATTR_IS_BIG,
     ATTR_KEYBOARD,
@@ -98,6 +101,7 @@
     PLATFORM_WEBHOOKS,
     SERVICE_ANSWER_CALLBACK_QUERY,
     SERVICE_DELETE_MESSAGE,
+    SERVICE_DOWNLOAD_FILE,
     SERVICE_EDIT_CAPTION,
     SERVICE_EDIT_MESSAGE,
     SERVICE_EDIT_MESSAGE_MEDIA,
@@ -331,6 +335,16 @@
     extra=vol.ALLOW_EXTRA,
 )
 
+SERVICE_SCHEMA_DOWNLOAD_FILE = vol.Schema(
+    {
+        vol.Optional(CONF_CONFIG_ENTRY_ID): cv.string,
+        vol.Required(ATTR_FILE_ID): cv.string,
+        vol.Optional(ATTR_DIRECTORY_PATH): cv.string,
+        vol.Optional(ATTR_FILE_NAME): cv.string,
+    },
+    extra=vol.ALLOW_EXTRA,
+)
+
 SERVICE_MAP = {
     SERVICE_SEND_MESSAGE: SERVICE_SCHEMA_SEND_MESSAGE,
     SERVICE_SEND_CHAT_ACTION: SERVICE_SCHEMA_SEND_CHAT_ACTION,
@@ -350,6 +364,7 @@
     SERVICE_DELETE_MESSAGE: SERVICE_SCHEMA_DELETE_MESSAGE,
     SERVICE_LEAVE_CHAT: SERVICE_SCHEMA_LEAVE_CHAT,
     SERVICE_SET_MESSAGE_REACTION: SERVICE_SCHEMA_SET_MESSAGE_REACTION,
+    SERVICE_DOWNLOAD_FILE: SERVICE_SCHEMA_DOWNLOAD_FILE,
 }
 
 
@@ -470,6 +485,8 @@ async def async_send_telegram_message(service: ServiceCall) -> ServiceResponse:
             await notify_service.set_message_reaction(context=service.context, **kwargs)
         elif msgtype == SERVICE_EDIT_MESSAGE_MEDIA:
             await notify_service.edit_message_media(context=service.context, **kwargs)
+        elif msgtype == SERVICE_DOWNLOAD_FILE:
+            return await notify_service.download_file(context=service.context, **kwargs)
         else:
             await notify_service.edit_message(
                 msgtype, context=service.context, **kwargs
@@ -515,6 +532,7 @@ async def async_send_telegram_message(service: ServiceCall) -> ServiceResponse:
             SERVICE_SEND_STICKER,
             SERVICE_SEND_LOCATION,
             SERVICE_SEND_POLL,
+            SERVICE_DOWNLOAD_FILE,
         ]:
             supports_response = SupportsResponse.OPTIONAL
 

homeassistant/components/telegram_bot/bot.py

@@ -5,6 +5,7 @@
 from collections.abc import Callable, Sequence
 import io
 import logging
+import os
 from ssl import SSLContext
 from types import MappingProxyType
 from typing import Any, cast
@@ -13,6 +14,7 @@
 from telegram import (
     Bot,
     CallbackQuery,
+    File,
     InlineKeyboardButton,
     InlineKeyboardMarkup,
     InputMedia,
@@ -45,6 +47,7 @@
 from homeassistant.core import Context, HomeAssistant
 from homeassistant.exceptions import HomeAssistantError, ServiceValidationError
 from homeassistant.helpers.dispatcher import async_dispatcher_send
+from homeassistant.util.json import JsonValueType
 from homeassistant.util.ssl import get_default_context, get_default_no_verify_context
 
 from .const import (
@@ -61,6 +64,7 @@
     ATTR_FILE_ID,
     ATTR_FILE_MIME_TYPE,
     ATTR_FILE_NAME,
+    ATTR_FILE_PATH,
     ATTR_FILE_SIZE,
     ATTR_FROM_FIRST,
     ATTR_FROM_LAST,
@@ -1033,6 +1037,54 @@ async def set_message_reaction(
             context=context,
         )
 
+    async def download_file(
+        self,
+        file_id: str,
+        directory_path: str | None = None,
+        file_name: str | None = None,
+        context: Context | None = None,
+        **kwargs: dict[str, Any],
+    ) -> dict[str, JsonValueType]:
+        """Download a file from Telegram."""
+        if not directory_path:
+            directory_path = self.hass.config.path(DOMAIN)
+        if not self.hass.config.is_allowed_path(directory_path):
+            raise ServiceValidationError(
+                "File path has not been configured in allowlist_external_dirs.",
+                translation_domain=DOMAIN,
+                translation_key="allowlist_external_dirs_error",
+            )
+        file: File = await self._send_msg(
+            self.bot.get_file,
+            "Error getting file",
+            None,
+            file_id=file_id,
+            context=context,
+        )
+        if not file_name:
+            file_name = (
+                os.path.basename(file.file_path) if file.file_path else "unknown_file"
+            )
+        custom_path = os.path.join(directory_path, file_name)
+        if not os.path.exists(directory_path):
+            _LOGGER.debug("directory %s does not exist, creating it", directory_path)
+
+            def mkdir() -> None:
+                os.makedirs(directory_path, exist_ok=True)
+
+            await self.hass.async_add_executor_job(mkdir)
+        _LOGGER.debug("Download file %s to %s", file_id, custom_path)
+        try:
+            await file.download_to_drive(custom_path=custom_path)
+        except Exception as exc:
+            _LOGGER.error("Error downloading file to %s: %s", custom_path, exc)
+            raise HomeAssistantError(
+                translation_domain=DOMAIN,
+                translation_key="action_failed",
+                translation_placeholders={"error": str(exc)},
+            ) from exc
+        return {ATTR_FILE_PATH: custom_path}
+
 
 def initialize_bot(hass: HomeAssistant, p_config: MappingProxyType[str, Any]) -> Bot:
     """Initialize telegram bot with proxy support."""

homeassistant/components/telegram_bot/const.py

@@ -50,6 +50,7 @@
 SERVICE_ANSWER_CALLBACK_QUERY = "answer_callback_query"
 SERVICE_DELETE_MESSAGE = "delete_message"
 SERVICE_LEAVE_CHAT = "leave_chat"
+SERVICE_DOWNLOAD_FILE = "download_file"
 
 SIGNAL_UPDATE_EVENT = "telegram_bot_update_event"
 EVENT_TELEGRAM_CALLBACK = "telegram_callback"
@@ -90,9 +91,11 @@
 ATTR_DATE = "date"
 ATTR_DISABLE_NOTIF = "disable_notification"
 ATTR_DISABLE_WEB_PREV = "disable_web_page_preview"
+ATTR_DIRECTORY_PATH = "directory_path"
 ATTR_EDITED_MSG = "edited_message"
 ATTR_FILE = "file"
 ATTR_FILE_ID = "file_id"
+ATTR_FILE_PATH = "file_path"
 ATTR_FILE_MIME_TYPE = "file_mime_type"
 ATTR_FILE_NAME = "file_name"
 ATTR_FILE_SIZE = "file_size"

homeassistant/components/telegram_bot/icons.json

@@ -13,6 +13,9 @@
     "delete_message": {
       "service": "mdi:delete"
     },
+    "download_file": {
+      "service": "mdi:paperclip"
+    },
     "edit_caption": {
       "service": "mdi:pencil"
     },

homeassistant/components/telegram_bot/services.yaml

@@ -977,3 +977,25 @@ set_message_reaction:
       required: false
       selector:
         boolean:
+
+download_file:
+  fields:
+    config_entry_id:
+      selector:
+        config_entry:
+          integration: telegram_bot
+    file_id:
+      required: true
+      example: "ABCD1234Efgh5678Ijkl90mnopQRStuvwx"
+      selector:
+        text:
+    directory_path:
+      required: false
+      default: "/config/telegram_bot"
+      selector:
+        text:
+    file_name:
+      required: false
+      example: "my_downloaded_file"
+      selector:
+        text:

homeassistant/components/telegram_bot/strings.json

@@ -327,6 +327,29 @@
       },
       "name": "Delete message"
     },
+    "download_file": {
+      "description": "Download the file to a local path.",
+      "fields": {
+        "config_entry_id": {
+          "description": "The config entry representing the Telegram bot to get the file.",
+          "name": "[%key:component::telegram_bot::services::send_message::fields::config_entry_id::name%]"
+        },
+        "directory_path": {
+          "description": "Local directory path to save the file to. defaults to the 'telegram_bot' directory within your Home Assistant configuration directory.",
+          "example": "/config/telegram_bot",
+          "name": "Directory path"
+        },
+        "file_id": {
+          "description": "ID of the file to get.",
+          "name": "File ID"
+        },
+        "file_name": {
+          "description": "Name to save the file as. If not provided, the original file name will be used.",
+          "name": "File name"
+        }
+      },
+      "name": "Download file"
+    },
     "edit_caption": {
       "description": "Edits the caption of a previously sent message.",
       "fields": {

tests/components/telegram_bot/test_telegram_bot.py

@@ -3,11 +3,19 @@
 import base64
 from datetime import datetime
 import io
+import os
 from typing import Any
 from unittest.mock import AsyncMock, MagicMock, mock_open, patch
 
 import pytest
-from telegram import Chat, InlineKeyboardButton, InlineKeyboardMarkup, Message, Update
+from telegram import (
+    Chat,
+    File,
+    InlineKeyboardButton,
+    InlineKeyboardMarkup,
+    Message,
+    Update,
+)
 from telegram.constants import ChatType, InputMediaType, ParseMode
 from telegram.error import (
     InvalidToken,
@@ -28,9 +36,12 @@
     ATTR_CAPTION,
     ATTR_CHAT_ACTION,
     ATTR_CHAT_ID,
+    ATTR_DIRECTORY_PATH,
     ATTR_DISABLE_NOTIF,
     ATTR_DISABLE_WEB_PREV,
     ATTR_FILE,
+    ATTR_FILE_ID,
+    ATTR_FILE_NAME,
     ATTR_KEYBOARD,
     ATTR_KEYBOARD_INLINE,
     ATTR_MEDIA_TYPE,
@@ -1473,3 +1484,122 @@ async def test_send_message_multi_target(
     await hass.async_block_till_done()
     assert response["chats"][0]["chat_id"] == 654321
     assert response["chats"][0]["message_id"] == 12345
+
+
+@pytest.mark.parametrize(
+    ("telegram_file_name", "schema_request", "expected_download_path"),
+    [
+        pytest.param(
+            "file_name1.jpg",
+            {ATTR_FILE_ID: "some-file-id-1"},
+            lambda hass: f"{hass.config.path(DOMAIN)}/file_name1.jpg",
+            id="no_custom_name_no_custom_dir",
+        ),
+        pytest.param(
+            "file_name2.jpg",
+            {ATTR_FILE_ID: "some-file-id-2", ATTR_FILE_NAME: "custom_name2.jpg"},
+            lambda hass: f"{hass.config.path(DOMAIN)}/custom_name2.jpg",
+            id="custom_name_no_custom_dir",
+        ),
+        pytest.param(
+            "file_name3.jpg",
+            {
+                ATTR_FILE_ID: "some-file-id-3",
+                ATTR_DIRECTORY_PATH: "/tmp/tempfolder3",  # noqa: S108
+            },
+            lambda hass: "/tmp/tempfolder3/file_name3.jpg",  # noqa: S108
+            id="no_custom_name_custom_dir",
+        ),
+        pytest.param(
+            "file_name4.jpg",
+            {
+                ATTR_FILE_ID: "some-file-id-4",
+                ATTR_FILE_NAME: "custom_name4.jpg",
+                ATTR_DIRECTORY_PATH: "/tmp/tempfolder4",  # noqa: S108
+            },
+            lambda hass: "/tmp/tempfolder4/custom_name4.jpg",  # noqa: S108
+            id="custom_name_custom_dir",
+        ),
+    ],
+)
+async def test_download_file(
+    hass: HomeAssistant,
+    mock_broadcast_config_entry: MockConfigEntry,
+    mock_external_calls: None,
+    telegram_file_name: str,
+    schema_request: dict[str, Any],
+    expected_download_path: Any,
+) -> None:
+    """Test download file."""
+    mock_broadcast_config_entry.add_to_hass(hass)
+    await hass.config_entries.async_setup(mock_broadcast_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    expected_path = expected_download_path(hass)
+    expected_directory = os.path.dirname(expected_path)
+    allowlist_dir = (
+        expected_directory
+        if os.path.exists(expected_directory)
+        else os.path.dirname(expected_directory)
+    )
+    hass.config.allowlist_external_dirs.add(allowlist_dir)
+
+    file_id = schema_request[ATTR_FILE_ID]
+    telegram_file = File(
+        file_id=file_id,
+        file_unique_id="file_unique_id",
+        file_path=f"file/path/{telegram_file_name}",
+    )
+
+    with (
+        patch(
+            "homeassistant.components.telegram_bot.bot.Bot.get_file",
+            AsyncMock(return_value=telegram_file),
+        ) as get_file_mock,
+        patch(
+            "telegram.File.download_to_drive",
+            AsyncMock(return_value=f"/custom/path/{telegram_file_name}"),
+        ) as download_to_drive_mock,
+    ):
+        await hass.services.async_call(
+            DOMAIN,
+            "download_file",
+            schema_request,
+            blocking=True,
+        )
+
+    await hass.async_block_till_done()
+    get_file_mock.assert_called_once_with(file_id=file_id)
+    download_to_drive_mock.assert_called_once_with(custom_path=expected_path)
+
+
+async def test_download_file_when_bot_failed_to_get_file(
+    hass: HomeAssistant,
+    mock_broadcast_config_entry: MockConfigEntry,
+    mock_external_calls: None,
+) -> None:
+    """Test download file when bot failed to get file."""
+    mock_broadcast_config_entry.add_to_hass(hass)
+    await hass.config_entries.async_setup(mock_broadcast_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    schema_request = {
+        ATTR_FILE_ID: "some-file-id",
+        ATTR_DIRECTORY_PATH: "/some/path",
+        ATTR_FILE_NAME: "custom_name.jpg",
+    }
+
+    with (
+        patch(
+            "homeassistant.components.telegram_bot.bot.Bot.get_file",
+            AsyncMock(side_effect=TelegramError("failed to get file")),
+        ),
+        pytest.raises(HomeAssistantError),
+    ):
+        await hass.services.async_call(
+            DOMAIN,
+            "download_file",
+            schema_request,
+            blocking=True,
+        )
+    await hass.async_block_till_done()