151 lines
5.7 KiB
Python
151 lines
5.7 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import time
|
|
import typing
|
|
from typing import cast
|
|
|
|
from icmplib import Host, SocketPermissionError, async_ping
|
|
|
|
from ..const import MAX_EXECUTOR_WORKERS
|
|
from ..entries import (
|
|
DashboardEntry,
|
|
EntryState,
|
|
EntryStateSource,
|
|
ReachableState,
|
|
bool_to_entry_state,
|
|
)
|
|
from ..util.itertools import chunked
|
|
|
|
if typing.TYPE_CHECKING:
|
|
from ..core import ESPHomeDashboard
|
|
|
|
|
|
_LOGGER = logging.getLogger(__name__)
|
|
|
|
GROUP_SIZE = int(MAX_EXECUTOR_WORKERS / 2)
|
|
|
|
DNS_FAILURE_STATE = EntryState(ReachableState.DNS_FAILURE, EntryStateSource.PING)
|
|
|
|
MIN_PING_INTERVAL = 5 # ensure we don't ping too often
|
|
|
|
|
|
class PingStatus:
|
|
def __init__(self, dashboard: ESPHomeDashboard) -> None:
|
|
"""Initialize the PingStatus class."""
|
|
super().__init__()
|
|
self._loop = asyncio.get_running_loop()
|
|
self.dashboard = dashboard
|
|
|
|
async def async_run(self) -> None:
|
|
"""Run the ping status."""
|
|
dashboard = self.dashboard
|
|
entries = dashboard.entries
|
|
privileged = await _can_use_icmp_lib_with_privilege()
|
|
if privileged is None:
|
|
_LOGGER.warning("Cannot use icmplib because privileges are insufficient")
|
|
return
|
|
|
|
while not dashboard.stop_event.is_set():
|
|
# Only ping if the dashboard is open
|
|
await dashboard.ping_request.wait()
|
|
dashboard.ping_request.clear()
|
|
iteration_start = time.monotonic()
|
|
current_entries = dashboard.entries.async_all()
|
|
to_ping: list[DashboardEntry] = []
|
|
|
|
for entry in current_entries:
|
|
if entry.address is None:
|
|
# No address or we already have a state from another source
|
|
# so no need to ping
|
|
continue
|
|
if (
|
|
entry.state.reachable is ReachableState.ONLINE
|
|
and entry.state.source
|
|
not in (EntryStateSource.PING, EntryStateSource.UNKNOWN)
|
|
):
|
|
# If we already have a state from another source and
|
|
# it's online, we don't need to ping
|
|
continue
|
|
to_ping.append(entry)
|
|
|
|
# Resolve DNS for all entries
|
|
entries_with_addresses: dict[DashboardEntry, list[str]] = {}
|
|
for ping_group in chunked(to_ping, GROUP_SIZE):
|
|
ping_group = cast(list[DashboardEntry], ping_group)
|
|
now_monotonic = time.monotonic()
|
|
dns_results = await asyncio.gather(
|
|
*(
|
|
dashboard.dns_cache.async_resolve(entry.address, now_monotonic)
|
|
for entry in ping_group
|
|
),
|
|
return_exceptions=True,
|
|
)
|
|
|
|
for entry, result in zip(ping_group, dns_results):
|
|
if isinstance(result, Exception):
|
|
# Only update state if its unknown or from ping
|
|
# so we don't mark it as offline if we have a state
|
|
# from mDNS or MQTT
|
|
entries.async_set_state_if_source(entry, DNS_FAILURE_STATE)
|
|
continue
|
|
if isinstance(result, BaseException):
|
|
raise result
|
|
entries_with_addresses[entry] = result
|
|
|
|
# Ping all entries with valid addresses
|
|
for ping_group in chunked(entries_with_addresses.items(), GROUP_SIZE):
|
|
entry_addresses = cast(tuple[DashboardEntry, list[str]], ping_group)
|
|
|
|
results = await asyncio.gather(
|
|
*(
|
|
async_ping(addresses[0], privileged=privileged)
|
|
for _, addresses in entry_addresses
|
|
),
|
|
return_exceptions=True,
|
|
)
|
|
|
|
for entry_addresses, result in zip(entry_addresses, results):
|
|
if isinstance(result, Exception):
|
|
ping_result = False
|
|
elif isinstance(result, BaseException):
|
|
raise result
|
|
else:
|
|
host: Host = result
|
|
ping_result = host.is_alive
|
|
entry: DashboardEntry = entry_addresses[0]
|
|
# If we can reach it via ping, we always set it
|
|
# online, however if we can't reach it via ping
|
|
# we only set it to offline if the state is unknown
|
|
# or from ping
|
|
entries.async_set_state_if_online_or_source(
|
|
entry,
|
|
bool_to_entry_state(ping_result, EntryStateSource.PING),
|
|
)
|
|
|
|
if not dashboard.stop_event.is_set():
|
|
iteration_duration = time.monotonic() - iteration_start
|
|
if iteration_duration < MIN_PING_INTERVAL:
|
|
await asyncio.sleep(MIN_PING_INTERVAL - iteration_duration)
|
|
|
|
|
|
async def _can_use_icmp_lib_with_privilege() -> None | bool:
|
|
"""Verify we can create a raw socket."""
|
|
try:
|
|
await async_ping("127.0.0.1", count=0, timeout=0, privileged=True)
|
|
except SocketPermissionError:
|
|
try:
|
|
await async_ping("127.0.0.1", count=0, timeout=0, privileged=False)
|
|
except SocketPermissionError:
|
|
_LOGGER.debug(
|
|
"Cannot use icmplib because privileges are insufficient to create the"
|
|
" socket"
|
|
)
|
|
return None
|
|
|
|
_LOGGER.debug("Using icmplib in privileged=False mode")
|
|
return False
|
|
|
|
_LOGGER.debug("Using icmplib in privileged=True mode")
|
|
return True
|