sarlink-portal-api/api/omada.py

import requests
from apibase.env import env


class Omada:
    def __init__(self):
        self.proxy_url = env("OMADA_PROXY_URL", default="")  # type: ignore
        self.api_key = env.str("OMADA_PROXY_API_KEY", default="")  # type: ignore
        self.site_id = env("OMADA_SITE_ID", default="")  # type: ignore
        self.group_id = env("OMADA_GROUP_ID", default="")  # type: ignore

        if not self.proxy_url:
            raise ValueError("OMADA_PROXY_URL is not set in the environment variables.")
        if not self.api_key:
            raise ValueError(
                "OMADA_PROXY_API_KEY is not set in the environment variables."
            )
        if not self.site_id:
            raise ValueError("OMADA_SITE_ID is not set in the environment variables.")
        if not self.group_id:
            raise ValueError("OMADA_GROUP_ID is not set in the environment variables.")

    def _get_existing_omada_devices(self):
        """
        Get existing Omada devices from the database.
        :return: List of existing device names.
        """
        try:
            url = f"{self.proxy_url}/9fd0cffa3475a74ae4e4d37de0d12414/api/v2/sites/{self.site_id}/setting/profiles/groups"
            response = requests.get(
                url,
                headers={"X-API-Key": str(self.api_key)},
            )
            print("Response: ", response.status_code)
            data = response.json()
            existing_devices = []
            if "result" in data and len(data["result"]["data"]) > 0:
                last_entry = data["result"]["data"][-1]
                print("Last Entry: ", last_entry)
                if "macAddressList" in last_entry:
                    existing_devices = last_entry["macAddressList"]
            print(existing_devices)
            return existing_devices
        except requests.RequestException as e:
            print(f"Error fetching existing devices: {e}")
            return []

    def add_new_devices_to_omada(self, new_devices: list[dict]):
        """
        Add new devices to Omada.
        :param new_devices: List of new device names to add.
        """
        try:
            PAYLOAD = {
                "name": "REGISTERED_DEVICES",
                "type": 2,
                "resource": 0,
                "ipList": None,
                "ipv6List": None,
                "macAddressList": None,
                "portList": None,
                "countryList": None,
                "portType": None,
                "portMaskList": None,
                "domainNamePort": None,
            }
            existing_devices = self._get_existing_omada_devices()
            PAYLOAD["macAddressList"] = existing_devices
            print("Payload with existing devices: ", PAYLOAD)
            for device in new_devices:
                print("Device in loop: ", device)
                PAYLOAD["macAddressList"].append(
                    {
                        "macAddress": device["mac"],
                        "name": device["name"],
                    }
                )
            print("New Payload: ", PAYLOAD)
            url = f"{self.proxy_url}/9fd0cffa3475a74ae4e4d37de0d12414/api/v2/sites/{self.site_id}/setting/profiles/groups/2/{self.group_id}"
            print(url)
            response = requests.patch(
                url,
                headers={"X-API-Key": str(self.api_key)},
                json=PAYLOAD,
            )
            print("Response: ", response.status_code)
            if response.status_code == 200:
                print("Devices successfully added.")
                print(response.json())
            else:
                print(f"Failed to add devices: {response.text}")
        except requests.RequestException as e:
            print(f"Error adding devices to Omada: {e}")

    def block_device(self, mac_address: str, operation: str = "block"):
        """
        Block a device in Omada.
        :param mac_address: The MAC address of the device to block.
        """
        try:
            url = f"{self.proxy_url}/9fd0cffa3475a74ae4e4d37de0d12414/api/v2/sites/{self.site_id}/cmd/clients/{mac_address}/{operation}"
            response = requests.post(
                url,
                headers={"X-API-Key": str(self.api_key)},
            )
            if response.status_code == 200:
                data = response.json()
                return data
            else:
                print(f"Failed to block device: {response.text}")
        except requests.RequestException as e:
            print(f"Error blocking device in Omada: {e}")