diff --git a/homeassistant/components/python_script/__init__.py b/homeassistant/components/python_script/__init__.py index 70e9c5b0d29a7..af7732780293d 100644 --- a/homeassistant/components/python_script/__init__.py +++ b/homeassistant/components/python_script/__init__.py @@ -1,5 +1,6 @@ """Component to allow running Python scripts.""" +from collections.abc import Mapping, Sequence import datetime import glob import logging @@ -7,6 +8,7 @@ import operator import os import time +import types from typing import Any from RestrictedPython import ( @@ -167,6 +169,20 @@ def python_script_service_handler(call: ServiceCall) -> ServiceResponse: } +def guarded_import( + name: str, + globals: Mapping[str, object] | None = None, + locals: Mapping[str, object] | None = None, + fromlist: Sequence[str] = (), + level: int = 0, +) -> types.ModuleType: + """Guard imports.""" + # Allow import of _strptime needed by datetime.datetime.strptime + if name == "_strptime": + return __import__(name, globals, locals, fromlist, level) + raise ScriptError(f"Not allowed to import {name}") + + def guarded_inplacevar(op: str, target: Any, operand: Any) -> Any: """Implement augmented-assign (+=, -=, etc.) operators for restricted code. @@ -232,6 +248,7 @@ def protected_getattr(obj, name, default=None): return getattr(obj, name, default) extra_builtins = { + "__import__": guarded_import, "datetime": datetime, "sorted": sorted, "time": TimeWrapper(), diff --git a/tests/components/python_script/test_init.py b/tests/components/python_script/test_init.py index c4dc00c448aae..2d151b4b81e6a 100644 --- a/tests/components/python_script/test_init.py +++ b/tests/components/python_script/test_init.py @@ -688,3 +688,27 @@ async def test_prohibited_augmented_assignment_operations( hass.async_add_executor_job(execute, hass, "aug_assign_prohibited.py", case, {}) await hass.async_block_till_done(wait_background_tasks=True) assert error in caplog.text + + +async def test_import_allow_strptime( + hass: HomeAssistant, caplog: pytest.LogCaptureFixture +) -> None: + """Test calling datetime.datetime.strptime works.""" + source = """ +test_date = datetime.datetime.strptime('2024-04-01', '%Y-%m-%d') +logger.info(f'Date {test_date}') + """ + hass.async_add_executor_job(execute, hass, "test.py", source, {}) + await hass.async_block_till_done(wait_background_tasks=True) + assert "Error executing script: Not allowed to import _strptime" not in caplog.text + assert "Date 2024-04-01 00:00:00" in caplog.text + + +async def test_no_other_imports_allowed( + hass: HomeAssistant, caplog: pytest.LogCaptureFixture +) -> None: + """Test imports are not allowed.""" + source = "import sys" + hass.async_add_executor_job(execute, hass, "test.py", source, {}) + await hass.async_block_till_done(wait_background_tasks=True) + assert "Error executing script: Not allowed to import sys" in caplog.text