From 414900fefb888eba50deb0901beaf871cfc0998f Mon Sep 17 00:00:00 2001 From: Paulus Schoutsen Date: Sun, 8 Oct 2017 23:51:32 -0700 Subject: [PATCH] Expose time module in Python Scripts (#9736) * Expose time module in Python Scripts * Make dt_util available in Python Scripts * Limit methods in time module * Add time.mktime * Limit access to datetime * Add warning to time.sleep * Lint --- homeassistant/components/python_script.py | 43 ++++++++++++++++-- tests/components/test_python_script.py | 53 ++++++++++++++++++++--- 2 files changed, 85 insertions(+), 11 deletions(-) diff --git a/homeassistant/components/python_script.py b/homeassistant/components/python_script.py index b33766d84db..6bf677b9645 100644 --- a/homeassistant/components/python_script.py +++ b/homeassistant/components/python_script.py @@ -1,8 +1,9 @@ """Component to allow running Python scripts.""" -import glob -import os -import logging import datetime +import glob +import logging +import os +import time import voluptuous as vol @@ -10,6 +11,7 @@ from homeassistant.const import SERVICE_RELOAD from homeassistant.exceptions import HomeAssistantError from homeassistant.loader import bind_hass from homeassistant.util import sanitize_filename +import homeassistant.util.dt as dt_util DOMAIN = 'python_script' REQUIREMENTS = ['restrictedpython==4.0a3'] @@ -25,6 +27,13 @@ ALLOWED_EVENTBUS = set(['fire']) ALLOWED_STATEMACHINE = set(['entity_ids', 'all', 'get', 'is_state', 'is_state_attr', 'remove', 'set']) ALLOWED_SERVICEREGISTRY = set(['services', 'has_service', 'call']) +ALLOWED_TIME = set(['sleep', 'strftime', 'strptime', 'gmtime', 'localtime', + 'ctime', 'time', 'mktime']) +ALLOWED_DATETIME = set(['date', 'time', 'datetime', 'timedelta', 'tzinfo']) +ALLOWED_DT_UTIL = set([ + 'utcnow', 'now', 'as_utc', 'as_timestamp', 'as_local', + 'utc_from_timestamp', 'start_of_local_day', 'parse_datetime', 'parse_date', + 'get_age']) class ScriptError(HomeAssistantError): @@ -111,7 +120,10 @@ def execute(hass, filename, source, data=None): elif (obj is hass and name not in ALLOWED_HASS or obj is hass.bus and name not in ALLOWED_EVENTBUS or obj is hass.states and name not in ALLOWED_STATEMACHINE or - obj is hass.services and name not in ALLOWED_SERVICEREGISTRY): + obj is hass.services and name not in ALLOWED_SERVICEREGISTRY or + obj is dt_util and name not in ALLOWED_DT_UTIL or + obj is datetime and name not in ALLOWED_DATETIME or + isinstance(obj, TimeWrapper) and name not in ALLOWED_TIME): raise ScriptError('Not allowed to access {}.{}'.format( obj.__class__.__name__, name)) @@ -120,6 +132,8 @@ def execute(hass, filename, source, data=None): builtins = safe_builtins.copy() builtins.update(utility_builtins) builtins['datetime'] = datetime + builtins['time'] = TimeWrapper() + builtins['dt_util'] = dt_util restricted_globals = { '__builtins__': builtins, '_print_': StubPrinter, @@ -159,3 +173,24 @@ class StubPrinter: # pylint: disable=no-self-use _LOGGER.warning( "Don't use print() inside scripts. Use logger.info() instead.") + + +class TimeWrapper: + """Wrapper of the time module.""" + + # Class variable, only going to warn once per Home Assistant run + warned = False + + # pylint: disable=no-self-use + def sleep(self, *args, **kwargs): + """Sleep method that warns once.""" + if not TimeWrapper.warned: + TimeWrapper.warned = True + _LOGGER.warning('Using time.sleep can reduce the performance of ' + 'Home Assistant') + + time.sleep(*args, **kwargs) + + def __getattr__(self, attr): + """Fetch an attribute from Time module.""" + return getattr(time, attr) diff --git a/tests/components/test_python_script.py b/tests/components/test_python_script.py index 660ed3c1b18..667d1849100 100644 --- a/tests/components/test_python_script.py +++ b/tests/components/test_python_script.py @@ -157,14 +157,17 @@ logger.info('Logging from inside script: %s %s' % (mydict["a"], mylist[2])) def test_accessing_forbidden_methods(hass, caplog): """Test compile error logs error.""" caplog.set_level(logging.ERROR) - source = """ -hass.stop() - """ - hass.async_add_job(execute, hass, 'test.py', source, {}) - yield from hass.async_block_till_done() - - assert "Not allowed to access HomeAssistant.stop" in caplog.text + for source, name in { + 'hass.stop()': 'HomeAssistant.stop', + 'dt_util.set_default_time_zone()': 'module.set_default_time_zone', + 'datetime.non_existing': 'module.non_existing', + 'time.tzset()': 'TimeWrapper.tzset', + }.items(): + caplog.records.clear() + hass.async_add_job(execute, hass, 'test.py', source, {}) + yield from hass.async_block_till_done() + assert "Not allowed to access {}".format(name) in caplog.text @asyncio.coroutine @@ -205,6 +208,26 @@ hass.states.set('hello.ab_list', '{}'.format(ab_list)) assert caplog.text == '' +@asyncio.coroutine +def test_exposed_modules(hass, caplog): + """Test datetime and time modules exposed.""" + caplog.set_level(logging.ERROR) + source = """ +hass.states.set('module.time', time.strftime('%Y', time.gmtime(521276400))) +hass.states.set('module.datetime', + datetime.timedelta(minutes=1).total_seconds()) +""" + + hass.async_add_job(execute, hass, 'test.py', source, {}) + yield from hass.async_block_till_done() + + assert hass.states.is_state('module.time', '1986') + assert hass.states.is_state('module.datetime', '60.0') + + # No errors logged = good + assert caplog.text == '' + + @asyncio.coroutine def test_reload(hass): """Test we can re-discover scripts.""" @@ -238,3 +261,19 @@ def test_reload(hass): assert hass.services.has_service('python_script', 'hello2') assert hass.services.has_service('python_script', 'world_beer') assert hass.services.has_service('python_script', 'reload') + + +@asyncio.coroutine +def test_sleep_warns_one(hass, caplog): + """Test time.sleep warns once.""" + caplog.set_level(logging.WARNING) + source = """ +time.sleep(2) +time.sleep(5) +""" + + with patch('homeassistant.components.python_script.time.sleep'): + hass.async_add_job(execute, hass, 'test.py', source, {}) + yield from hass.async_block_till_done() + + assert caplog.text.count('time.sleep') == 1