unifi-traffic-routes-domain-ip/main.py
Ian Fijolek 95bc1528ec
All checks were successful
continuous-integration/drone/push Build is passing
continuous-integration/drone/tag Build is passing
Fix for newer aiounifi
2024-09-19 15:58:13 -07:00

120 lines
3.7 KiB
Python

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.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())