From f2308256514c6c0611802a652ca1953723011f35 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Sun, 26 Nov 2023 11:04:17 -0600
Subject: [PATCH] Add discover cli tool (#732)

---
 .coveragerc               |  1 +
 README.rst                |  6 +++
 aioesphomeapi/discover.py | 80 +++++++++++++++++++++++++++++++++++++++
 setup.py                  |  3 +-
 4 files changed, 89 insertions(+), 1 deletion(-)
 create mode 100644 aioesphomeapi/discover.py

diff --git a/.coveragerc b/.coveragerc
index dba3dc7..acb52e6 100644
--- a/.coveragerc
+++ b/.coveragerc
@@ -5,6 +5,7 @@ omit =
     aioesphomeapi/api_options_pb2.py
     aioesphomeapi/api_pb2.py
     aioesphomeapi/log_reader.py
+    aioesphomeapi/discover.py
     bench/*.py
 
 [report]
diff --git a/README.rst b/README.rst
index adbe623..cd2d47c 100644
--- a/README.rst
+++ b/README.rst
@@ -135,6 +135,12 @@ A cli tool is also available for watching logs:
 
    aioesphomeapi-logs --help
 
+A cli tool is also available to discover devices:
+
+.. code:: bash
+
+   aioesphomeapi-discover
+
 License
 -------
 
diff --git a/aioesphomeapi/discover.py b/aioesphomeapi/discover.py
new file mode 100644
index 0000000..52ba87f
--- /dev/null
+++ b/aioesphomeapi/discover.py
@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+# Helper script and aioesphomeapi to discover api devices
+import asyncio
+import logging
+import sys
+
+from zeroconf import IPVersion, ServiceStateChange, Zeroconf
+from zeroconf.asyncio import AsyncServiceBrowser, AsyncServiceInfo, AsyncZeroconf
+
+FORMAT = "{: <7}|{: <32}|{: <15}|{: <12}|{: <16}|{: <10}|{: <32}"
+COLUMN_NAMES = ("Status", "Name", "Address", "MAC", "Version", "Platform", "Board")
+
+
+def decode_bytes_or_none(data: str | bytes | None) -> str | None:
+    """Decode bytes or return None."""
+    if data is None:
+        return None
+    if isinstance(data, bytes):
+        return data.decode()
+    return data
+
+
+def async_service_update(
+    zeroconf: Zeroconf,
+    service_type: str,
+    name: str,
+    state_change: ServiceStateChange,
+) -> None:
+    """Service state changed."""
+    short_name = name.partition(".")[0]
+    if state_change is ServiceStateChange.Removed:
+        state = "OFFLINE"
+    else:
+        state = "ONLINE"
+    info = AsyncServiceInfo(service_type, name)
+    info.load_from_cache(zeroconf)
+    properties = info.properties
+    mac = decode_bytes_or_none(properties.get(b"mac"))
+    version = decode_bytes_or_none(properties.get(b"version"))
+    platform = decode_bytes_or_none(properties.get(b"platform"))
+    board = decode_bytes_or_none(properties.get(b"board"))
+    address = None
+    if addresses := info.ip_addresses_by_version(IPVersion.V4Only):
+        address = str(addresses[0])
+
+    print(FORMAT.format(state, short_name, address, mac, version, platform, board))
+
+
+async def main() -> None:
+    logging.basicConfig(
+        format="%(asctime)s.%(msecs)03d %(levelname)-8s %(message)s",
+        level=logging.INFO,
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+    aiozc = AsyncZeroconf()
+    browser = AsyncServiceBrowser(
+        aiozc.zeroconf, "_esphomelib._tcp.local.", handlers=[async_service_update]
+    )
+    print(FORMAT.format(*COLUMN_NAMES))
+    print("-" * 120)
+
+    try:
+        await asyncio.Event().wait()
+    finally:
+        await browser.async_cancel()
+        await aiozc.async_close()
+
+
+def cli_entry_point() -> None:
+    """Run the CLI."""
+    try:
+        asyncio.run(main())
+    except KeyboardInterrupt:
+        pass
+
+
+if __name__ == "__main__":
+    cli_entry_point()
+    sys.exit(0)
diff --git a/setup.py b/setup.py
index 124e64f..e99f302 100644
--- a/setup.py
+++ b/setup.py
@@ -61,7 +61,8 @@ setup_kwargs = {
     "test_suite": "tests",
     "entry_points": {
         "console_scripts": [
-            "aioesphomeapi-logs=aioesphomeapi.log_reader:cli_entry_point"
+            "aioesphomeapi-logs=aioesphomeapi.log_reader:cli_entry_point",
+            "aioesphomeapi-discover=aioesphomeapi.discover:cli_entry_point",
         ],
     },
 }