From 5f7d710404ee31fe6f77b2b0936348a3cc74db04 Mon Sep 17 00:00:00 2001 From: Ian Fijolek Date: Mon, 20 Nov 2023 09:31:58 -0800 Subject: [PATCH] Add the good stuff --- .drone.yml | 46 ++++++++++++++++++ Dockerfile | 9 ++++ README.md | 15 +++++- main.py | 122 +++++++++++++++++++++++++++++++++++++++++++++++ requirements.txt | 2 + 5 files changed, 193 insertions(+), 1 deletion(-) create mode 100644 .drone.yml create mode 100644 Dockerfile create mode 100644 main.py create mode 100644 requirements.txt diff --git a/.drone.yml b/.drone.yml new file mode 100644 index 0000000..9d83ac0 --- /dev/null +++ b/.drone.yml @@ -0,0 +1,46 @@ +--- +kind: pipeline +name: publish + +trigger: + event: + - push + - tag + refs: + - refs/heads/master + - refs/tags/v* + +steps: + - name: Build and publish + image: woodpeckerci/plugin-docker-buildx + settings: + repo: iamthefij/unifi-traffic-routes + auto_tag: true + username: + from_secret: docker_username + password: + from_secret: docker_password + +--- +kind: pipeline +name: notify + +depends_on: + - publish + +trigger: + status: + - failure + +steps: + + - name: notify + image: drillster/drone-email + settings: + host: + from_secret: SMTP_HOST # pragma: whitelist secret + username: + from_secret: SMTP_USER # pragma: whitelist secret + password: + from_secret: SMTP_PASS # pragma: whitelist secret + from: drone@iamthefij.com diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..53425c4 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,9 @@ +FROM python:3 + +COPY requirements.txt /src/ +RUN pip install --no-cache-dir -r /src/requirements.txt + +COPY main.py /src/ + +ENTRYPOINT ["/usr/local/bin/python"] +CMD ["/src/main.py"] diff --git a/README.md b/README.md index a9e49fc..31b0cb4 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,16 @@ # unfifi-traffic-routes-domain-ip -Look up and set traffic route IPs based on domains for clients not using the UniFi DNS \ No newline at end of file +Look up and set traffic route IPs based on domains for clients not using the UniFi DNS + +## Configuration + +Set the following as environment variables + +| Variable | Description | Default | +| -------- | ----------- | ------- | +| `UNIFI_HOST` | Unifi network hostname | `192.168.1.1` | +| `UNIFI_PORT` | Unifi network port number | `443` | +| `UNIFI_USER` | Unifi username | | +| `UNIFI_PASS` | Unifi pass | | + +The user must have access to Manage the Network application to update the Traffic Routes. diff --git a/main.py b/main.py new file mode 100644 index 0000000..4285e8f --- /dev/null +++ b/main.py @@ -0,0 +1,122 @@ +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_PASSWORD") + + +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()) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..8bc7dd3 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,2 @@ +# aiounifi +git+https://github.com/ViViDboarder/aiounifi.git@traffic-routes