discover.py

#!/usr/bin/env python3
import argparse
import json
import socket
import time
from typing import Any, Dict, List, Optional

import requests
from zeroconf import IPVersion, ServiceBrowser, ServiceInfo, Zeroconf

DEFAULT_SERVICE = "_opensentry._tcp.local."


def _decode_props(props: Dict[bytes, bytes]) -> Dict[str, str]:
    out: Dict[str, str] = {}
    for k, v in (props or {}).items():
        try:
            ks = k.decode() if isinstance(k, (bytes, bytearray)) else str(k)
        except Exception:
            ks = str(k)
        try:
            vs = v.decode() if isinstance(v, (bytes, bytearray)) else str(v)
        except Exception:
            vs = str(v)
        out[ks] = vs
    return out


class _Collector:
    def __init__(self, zc: Zeroconf, service_type: str):
        self.zc = zc
        self.service_type = service_type
        self.names: List[str] = []
        self._infos: Dict[str, ServiceInfo] = {}

    def add_service(self, zc: Zeroconf, type_: str, name: str) -> None:  # type: ignore[override]
        info = zc.get_service_info(type_, name, timeout=1000)
        if name not in self.names:
            self.names.append(name)
        if info:
            self._infos[name] = info

    def update_service(self, zc: Zeroconf, type_: str, name: str) -> None:  # type: ignore[override]
        info = zc.get_service_info(type_, name, timeout=1000)
        if info:
            self._infos[name] = info

    def remove_service(self, zc: Zeroconf, type_: str, name: str) -> None:  # type: ignore[override]
        self._infos.pop(name, None)
        if name in self.names:
            self.names.remove(name)

    def infos(self) -> List[ServiceInfo]:
        # Resolve unresolved names one more time
        for n in list(self.names):
            if n not in self._infos:
                info = self.zc.get_service_info(self.service_type, n, timeout=1000)
                if info:
                    self._infos[n] = info
        return list(self._infos.values())


def browse(service: str, timeout: float) -> List[ServiceInfo]:
    zc = Zeroconf(ip_version=IPVersion.V4Only)
    try:
        collector = _Collector(zc, service)
        browser = ServiceBrowser(zc, service, listener=collector)  # noqa: F841
        time.sleep(max(0.5, timeout))
        return collector.infos()
    finally:
        try:
            zc.close()
        except Exception:
            pass


def fetch_status(ip: str, port: int, token: Optional[str], timeout: float = 1.5) -> Optional[Dict[str, Any]]:
    url = f"http://{ip}:{port}/status"
    headers = {}
    if token:
        headers["Authorization"] = f"Bearer {token}"
    try:
        r = requests.get(url, headers=headers, timeout=timeout)
        if r.status_code == 200:
            return r.json()
    except Exception:
        return None
    return None


def to_result(info: ServiceInfo, with_status: bool, token: Optional[str]) -> Dict[str, Any]:
    addrs = []
    for a in info.addresses or []:
        try:
            addrs.append(socket.inet_ntoa(a))
        except Exception:
            pass
    props = _decode_props(getattr(info, "properties", {}) or {})
    port = int(getattr(info, "port", 0) or 0)
    name = getattr(info, "name", "")
    server = getattr(info, "server", "")

    entry: Dict[str, Any] = {
        "service": info.type if hasattr(info, "type") else DEFAULT_SERVICE,
        "name": name,
        "host": server,
        "ips": addrs,
        "port": port,
        "txt": props,
    }

    # Convenience fields
    if addrs:
        entry["ip"] = addrs[0]
        entry["url"] = f"http://{addrs[0]}:{port}/"

    if with_status and addrs and port:
        # Attempt status against first address
        st = fetch_status(addrs[0], port, token)
        if st is not None:
            entry["status"] = st
        else:
            entry["status_error"] = "unreachable or unauthorized"

    return entry


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Discover OpenSentry devices via mDNS")
    p.add_argument("--timeout", type=float, default=3.0, help="mDNS browse duration in seconds")
    p.add_argument("--service", default=DEFAULT_SERVICE, help="mDNS service type to browse")
    p.add_argument("--json", action="store_true", help="JSON output only")
    p.add_argument("--status", action="store_true", help="Fetch /status for each device")
    p.add_argument("--token", default=None, help="Bearer token for /status if required")
    return p.parse_args()


def main() -> int:
    args = parse_args()
    infos = browse(args.service, args.timeout)
    results = [to_result(i, args.status, args.token) for i in infos]

    if args.json:
        print(json.dumps(results, indent=2))
    else:
        if not results:
            print("No OpenSentry devices discovered via mDNS.")
            return 1
        print("Discovered OpenSentry devices:")
        for r in results:
            line = f"- {r.get('name','')} @ {r.get('ip','?')}:{r.get('port','?')}"
            txt = r.get("txt", {})
            if "id" in txt:
                line += f"  id={txt['id']}"
            if "auth" in txt:
                line += f"  auth={txt['auth']}"
            print(line)
            if args.status and "status" in r:
                st = r["status"]
                print(f"  status: version={st.get('version')} caps={st.get('caps')}")
            elif args.status and "status_error" in r:
                print(f"  status: {r['status_error']}")

    return 0


if __name__ == "__main__":
    raise SystemExit(main())