From 9688635f44a3d3fd73e0b8e4864b206836e32ae5 Mon Sep 17 00:00:00 2001 From: i701 Date: Sun, 22 Jun 2025 22:55:28 +0500 Subject: [PATCH] Refactor Omada integration: encapsulate API calls in Omada class, update device management tasks, and enhance device creation tests --- .python-version | 1 + api/omada.py | 93 +++++++++++++++++++++++++++++++++++++++++++ api/tasks.py | 101 ++++++----------------------------------------- devices/tests.py | 22 +++++++++-- devices/views.py | 1 - 5 files changed, 126 insertions(+), 92 deletions(-) create mode 100644 .python-version create mode 100644 api/omada.py diff --git a/.python-version b/.python-version new file mode 100644 index 0000000..d9506ce --- /dev/null +++ b/.python-version @@ -0,0 +1 @@ +3.12.5 diff --git a/api/omada.py b/api/omada.py new file mode 100644 index 0000000..f609270 --- /dev/null +++ b/api/omada.py @@ -0,0 +1,93 @@ +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}") diff --git a/api/tasks.py b/api/tasks.py index 559ea4f..f0a940c 100644 --- a/api/tasks.py +++ b/api/tasks.py @@ -2,23 +2,19 @@ from django.shortcuts import get_object_or_404 from api.models import User from devices.models import Device from api.notifications import send_sms -import requests -from apibase.env import env, BASE_DIR import os import logging from celery import shared_task from django.utils import timezone from api.notifications import send_clean_telegram_markdown +from api.omada import Omada +from apibase.env import env, BASE_DIR logger = logging.getLogger(__name__) - env.read_env(os.path.join(BASE_DIR, ".env")) -PERSON_VERIFY_BASE_URL = env.str("PERSON_VERIFY_BASE_URL", default="") # type: ignore -OMADA_PROXY_API_KEY = env.str("OMADA_PROXY_API_KEY", default="") # type: ignore -OMADA_PROXY_URL = env("OMADA_PROXY_URL", default="") # type: ignore -OMADA_SITE_ID = env("OMADA_SITE_ID", default="") # type: ignore -OMADA_GROUP_ID = env("OMADA_GROUP_ID", default="") # type: ignore + +omada_client = Omada() @shared_task @@ -61,93 +57,19 @@ def deactivate_expired_devices(): def get_existing_omada_devices(): """ - Get existing Omada devices from the database. + Get existing Omada devices from Omada API via Omada class. :return: List of existing device names. """ - if not OMADA_PROXY_URL: - raise ValueError( - "OMADA_PROXY_URL is not set. Please set it in your environment variables." - ) - try: - response = requests.get( - f"{OMADA_PROXY_URL}/9fd0cffa3475a74ae4e4d37de0d12414/api/v2/sites/66dcddb804aa0d2978cf145f/setting/profiles/groups", - headers={"X-API-Key": str(OMADA_PROXY_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 [] + return omada_client.get_existing_omada_devices() @shared_task def add_new_devices_to_omada(new_devices: list[dict]): """ - Add new devices to Omada. + Add new devices to Omada via Omada class. :param new_devices: List of new device names to add. """ - if not OMADA_SITE_ID: - raise ValueError( - "OMADA_SITE_ID is not set. Please set it in your environment variables." - ) - if not OMADA_GROUP_ID: - raise ValueError( - "OMADA_GROUP_ID is not set. Please set it in your environment variables." - ) - if not OMADA_PROXY_URL: - raise ValueError( - "OMADA_PROXY_URL is not set. Please set it in your environment variables." - ) - 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 = 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) - print( - f"{OMADA_PROXY_URL}/9fd0cffa3475a74ae4e4d37de0d12414/api/v2/sites/{OMADA_SITE_ID}/setting/profiles/groups/2/{OMADA_GROUP_ID}" - ) - response = requests.patch( - f"{OMADA_PROXY_URL}/9fd0cffa3475a74ae4e4d37de0d12414/api/v2/sites/{OMADA_SITE_ID}/setting/profiles/groups/2/{OMADA_GROUP_ID}", - headers={"X-API-Key": str(OMADA_PROXY_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: {e}") + omada_client.add_new_devices_to_omada(new_devices) def verify_user_with_person_api_task(user_id: int): @@ -155,7 +77,6 @@ def verify_user_with_person_api_task(user_id: int): Verify the user with the Person API. :param user_id: The ID of the user to verify. """ - user = get_object_or_404(User, id=user_id) # Call the Person API to verify the user @@ -171,10 +92,14 @@ def verify_user_with_person_api_task(user_id: int): """ logger.info(verification_failed_message) + PERSON_VERIFY_BASE_URL = env.str("PERSON_VERIFY_BASE_URL", default="") # type: ignore + if not PERSON_VERIFY_BASE_URL: raise ValueError( "PERSON_VERIFY_BASE_URL is not set in the environment variables." ) + import requests + response = requests.get(f"{PERSON_VERIFY_BASE_URL}/api/person/{user.id_card}") if response.status_code == 200: data = response.json() @@ -250,4 +175,4 @@ def verify_user_with_person_api_task(user_id: int): else: # Handle the error case print(f"Error verifying user: {response.status_code} - {response.text}") - return False + return diff --git a/devices/tests.py b/devices/tests.py index efc28f6..8a9aa9d 100644 --- a/devices/tests.py +++ b/devices/tests.py @@ -35,11 +35,18 @@ class DeviceAPITestCase(APITestCase): response = self.client.get("/api/devices/") self.assertEqual(response.status_code, status.HTTP_200_OK) - def test_create_device(self): + def test_create_device_with_no_vendor_mac(self): self.client.credentials(HTTP_AUTHORIZATION=f"Token {self.token}") data = {"name": "New Device", "mac": "11:22:33:44:55:66"} response = self.client.post("/api/devices/", data) + self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) + + def test_create_device_with_vendor_mac(self): + self.client.credentials(HTTP_AUTHORIZATION=f"Token {self.token}") + data = {"name": "New Device", "mac": "64:16:7F:21:51:A5"} + response = self.client.post("/api/devices/", data) self.assertEqual(response.status_code, status.HTTP_201_CREATED) + self.assertIn("New Device", response.data["name"]) def test_create_device_invalid_mac(self): self.client.credentials(HTTP_AUTHORIZATION=f"Token {self.token}") @@ -48,6 +55,15 @@ class DeviceAPITestCase(APITestCase): self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) self.assertIn("Invalid mac address.", response.data["message"]) + def test_mac_address_already_exists(self): + self.client.credentials(HTTP_AUTHORIZATION=f"Token {self.token}") + data = {"name": "Existing Device", "mac": self.device.mac} + response = self.client.post("/api/devices/", data) + self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) + self.assertIn( + "Device with this mac address already exists.", response.data["message"] + ) + def test_update_device(self): self.client.credentials(HTTP_AUTHORIZATION=f"Token {self.token}") data = {"name": "Updated Device"} @@ -59,14 +75,14 @@ class DeviceAPITestCase(APITestCase): def test_block_device(self): self.client.credentials(HTTP_AUTHORIZATION=f"Token {self.token}") data = { - "blocked": True, # also use Python boolean here, not a string + "blocked": True, "reason_for_blocking": "hello", "blocked_by": "ADMIN", } response = self.client.put( f"/api/devices/{self.device.pk}/block/", data, - format="json", # ✅ this is crucial! + format="json", ) self.assertEqual(response.status_code, status.HTTP_200_OK) self.device.refresh_from_db() diff --git a/devices/views.py b/devices/views.py index aec9cde..19d1cf5 100644 --- a/devices/views.py +++ b/devices/views.py @@ -63,7 +63,6 @@ class DeviceListCreateAPIView( return Response({"message": "MAC address vendor not found."}, status=400) mac = re.sub(NORMALIZE_MAC_REGEX, "-", mac).upper() - request.data["mac"] = mac return super().create(request, *args, **kwargs)