unifi-traffic-routes-domain-ip/main.py

import asyncio
import logging
import os
import socket
from asyncio.timeouts import timeout

import aiohttp
import aiounifi
from aiounifi.controller import Controller
from aiounifi.models.configuration import Configuration
from aiounifi.models.traffic_route import IPAddress
from aiounifi.models.traffic_route import MatchingTarget

HOST = os.getenv("UNIFI_HOST", "192.168.1.1")
PORT = int(os.getenv("UNIFI_PORT", 443))
USER = os.getenv("UNIFI_USER")
PASS = os.getenv("UNIFI_PASS")


LOGGER = logging.getLogger(__name__)


def get_configuration(session: aiohttp.ClientSession) -> Configuration:
    return Configuration(
        session,
        HOST,
        username=USER,
        password=PASS,
        port=PORT,
    )


async def get_controller(config: Configuration) -> Controller | None:
    controller = Controller(config)

    try:
        async with timeout(10):
            await controller.login()
        return controller
    except aiounifi.LoginRequired:
        LOGGER.error("Connected to UniFi at %s:%s but couldn't log in", config.host, config.port)

    except aiounifi.Unauthorized:
        LOGGER.error("Connected to UniFi at %s:%s but not registered", config.host, config.port)

    except (asyncio.TimeoutError, aiounifi.RequestError):
        LOGGER.exception("Error connecting to the UniFi controller at %s:%s", config.host, config.port)

    except aiounifi.AiounifiException:
        LOGGER.exception("Unknown UniFi communication error occurred")

    return None


async def main():
    session = aiohttp.ClientSession(cookie_jar=aiohttp.CookieJar(unsafe=True))
    config = get_configuration(session)
    controller = await get_controller(config)
    if not controller:
        LOGGER.error("Couldn't connect")
        await session.close()
        return

    try:
        await controller.initialize()

        print("Initialized!!")

        await controller.traffic_routes.update()
        for item in controller.traffic_routes.values():
            if item.domains and item.matching_target == MatchingTarget.IP:
                print(item.description)

                # Look up unique ip addresses
                addresses = set()
                for domain in item.domains:
                    print(domain["domain"])
                    addresses.update(
                        result[4][0]
                        for result in socket.getaddrinfo(
                            domain["domain"],
                            80,
                            type=socket.SOCK_STREAM,
                            proto=socket.IPPROTO_IP,
                        )
                    )

                # Sort addresses by type and value
                sorted_addresses = list(
                    sorted(
                        addresses, key=lambda a: f"{'ip4' if '.' in a else 'ip6'}-{a}"
                    )
                )
                print(sorted_addresses)

                # Build addresses objects
                ip_addresses = [
                    IPAddress(
                        ip_or_subnet=a,
                        ip_version="v4" if "." in a else "v6",
                        port_ranges=[],
                        ports=[],
                    )
                    for a in sorted_addresses
                    if a not in ("0.0.0.0", "::")
                ]

                # Check for change
                if ip_addresses == item.raw["ip_addresses"]:
                    print("already up to date")
                    continue

                # Update while preserving current state
                item.raw["ip_addresses"] = ip_addresses
                result = await controller.traffic_routes.save(item)
                print(result)

    finally:
        await session.close()


if __name__ == "__main__":
    asyncio.run(main())