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

123 lines
3.7 KiB
Python
Raw Normal View History

2023-11-20 17:31:58 +00:00
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))
USERNAME = os.getenv("UNIFI_USER")
PASSWORD = os.getenv("UNIFI_PASS")
2023-11-20 17:31:58 +00:00
LOGGER = logging.getLogger(__name__)
def get_configuration(session: aiohttp.ClientSession) -> Configuration:
return Configuration(
session,
HOST,
username=USERNAME,
password=PASSWORD,
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
]
# 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())