124 lines
3.7 KiB
Python
124 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.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())
|