From 5b0ee473b6ec7d7f636e814a6c0c661537317748 Mon Sep 17 00:00:00 2001
From: Paulus Schoutsen <balloob@gmail.com>
Date: Tue, 23 Apr 2019 03:46:23 -0700
Subject: [PATCH] Add get_states faster (#23315)

---
 homeassistant/auth/permissions/__init__.py    | 13 ++++++++++++
 homeassistant/auth/permissions/util.py        | 14 +++++++++++++
 .../components/websocket_api/commands.py      | 13 +++++++-----
 tests/auth/permissions/test_util.py           | 21 +++++++++++++++++++
 4 files changed, 56 insertions(+), 5 deletions(-)
 create mode 100644 tests/auth/permissions/test_util.py

diff --git a/homeassistant/auth/permissions/__init__.py b/homeassistant/auth/permissions/__init__.py
index 63e76dd2496..0079f11447b 100644
--- a/homeassistant/auth/permissions/__init__.py
+++ b/homeassistant/auth/permissions/__init__.py
@@ -11,6 +11,7 @@ from .models import PermissionLookup
 from .types import PolicyType
 from .entities import ENTITY_POLICY_SCHEMA, compile_entities
 from .merge import merge_policies  # noqa
+from .util import test_all
 
 
 POLICY_SCHEMA = vol.Schema({
@@ -29,6 +30,10 @@ class AbstractPermissions:
         """Return a function that can test entity access."""
         raise NotImplementedError
 
+    def access_all_entities(self, key: str) -> bool:
+        """Check if we have a certain access to all entities."""
+        raise NotImplementedError
+
     def check_entity(self, entity_id: str, key: str) -> bool:
         """Check if we can access entity."""
         entity_func = self._cached_entity_func
@@ -48,6 +53,10 @@ class PolicyPermissions(AbstractPermissions):
         self._policy = policy
         self._perm_lookup = perm_lookup
 
+    def access_all_entities(self, key: str) -> bool:
+        """Check if we have a certain access to all entities."""
+        return test_all(self._policy.get(CAT_ENTITIES), key)
+
     def _entity_func(self) -> Callable[[str, str], bool]:
         """Return a function that can test entity access."""
         return compile_entities(self._policy.get(CAT_ENTITIES),
@@ -65,6 +74,10 @@ class _OwnerPermissions(AbstractPermissions):
 
     # pylint: disable=no-self-use
 
+    def access_all_entities(self, key: str) -> bool:
+        """Check if we have a certain access to all entities."""
+        return True
+
     def _entity_func(self) -> Callable[[str, str], bool]:
         """Return a function that can test entity access."""
         return lambda entity_id, key: True
diff --git a/homeassistant/auth/permissions/util.py b/homeassistant/auth/permissions/util.py
index d2d259fb32e..0d334c4a3ba 100644
--- a/homeassistant/auth/permissions/util.py
+++ b/homeassistant/auth/permissions/util.py
@@ -3,6 +3,7 @@ from functools import wraps
 
 from typing import Callable, Dict, List, Optional, Union, cast  # noqa: F401
 
+from .const import SUBCAT_ALL
 from .models import PermissionLookup
 from .types import CategoryType, SubCategoryDict, ValueType
 
@@ -96,3 +97,16 @@ def _gen_dict_test_func(
         return schema.get(key)
 
     return test_value
+
+
+def test_all(policy: CategoryType, key: str) -> bool:
+    """Test if a policy has an ALL access for a specific key."""
+    if not isinstance(policy, dict):
+        return bool(policy)
+
+    all_policy = policy.get(SUBCAT_ALL)
+
+    if not isinstance(all_policy, dict):
+        return bool(all_policy)
+
+    return all_policy.get(key, False)
diff --git a/homeassistant/components/websocket_api/commands.py b/homeassistant/components/websocket_api/commands.py
index d9834758c80..84178beef8b 100644
--- a/homeassistant/components/websocket_api/commands.py
+++ b/homeassistant/components/websocket_api/commands.py
@@ -142,11 +142,14 @@ def handle_get_states(hass, connection, msg):
 
     Async friendly.
     """
-    entity_perm = connection.user.permissions.check_entity
-    states = [
-        state for state in hass.states.async_all()
-        if entity_perm(state.entity_id, 'read')
-    ]
+    if connection.user.permissions.access_all_entities('read'):
+        states = hass.states.async_all()
+    else:
+        entity_perm = connection.user.permissions.check_entity
+        states = [
+            state for state in hass.states.async_all()
+            if entity_perm(state.entity_id, 'read')
+        ]
 
     connection.send_message(messages.result_message(
         msg['id'], states))
diff --git a/tests/auth/permissions/test_util.py b/tests/auth/permissions/test_util.py
new file mode 100644
index 00000000000..1a339208f4d
--- /dev/null
+++ b/tests/auth/permissions/test_util.py
@@ -0,0 +1,21 @@
+"""Test the permission utils."""
+
+from homeassistant.auth.permissions import util
+
+
+def test_test_all():
+    """Test if we can test the all group."""
+    for val in (
+            None,
+            {},
+            {'all': None},
+            {'all': {}},
+    ):
+        assert util.test_all(val, 'read') is False
+
+    for val in (
+            True,
+            {'all': True},
+            {'all': {'read': True}},
+    ):
+        assert util.test_all(val, 'read') is True