Fix strptime in python_script

homeassistant/components/python_script/__init__.py

@@ -1,12 +1,14 @@
 """Component to allow running Python scripts."""
 
+from collections.abc import Mapping, Sequence
 import datetime
 import glob
 import logging
 from numbers import Number
 import operator
 import os
 import time
+import types
 from typing import Any
 
 from RestrictedPython import (
@@ -167,6 +169,19 @@ 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. Only allow datetime to import _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 +247,7 @@ def protected_getattr(obj, name, default=None):
         return getattr(obj, name, default)
 
     extra_builtins = {
+        "__import__": guarded_import,
         "datetime": datetime,
         "sorted": sorted,
         "time": TimeWrapper(),

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 that prohibited augmented assignment operations raise an error."""
+    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 that prohibited augmented assignment operations raise an error."""
+    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