[changes] Tasks Refactoring, documentation update

Signed-off-by: DragnEmperor <[email protected]>

Author: DragnEmperor

docs/user/estimated-location.rst

@@ -16,14 +16,14 @@ Estimated Location
 Overview
 --------
 
-The Estimated Location feature automatically creates or updates a device’s
-location based on latitude and longitude information retrieved from the
-WHOIS Lookup feature.
+This feature automatically creates or updates a device’s location based on
+latitude and longitude information retrieved from the WHOIS Lookup
+feature.
 
 Trigger Conditions
 ------------------
 
-Estimated Location is triggered when:
+This feature is triggered when:
 
 - A **fresh WHOIS lookup** is performed for a device.
 - Or when a WHOIS record already exists for the device’s IP **and**:

docs/user/whois.rst

@@ -15,10 +15,10 @@ WHOIS Lookup
 Overview
 --------
 
-The WHOIS Lookup feature displays information about the public IP address
-used by devices to communicate with OpenWISP (via the ``last_ip`` field).
-It helps identify the geographic location and ISP associated with the IP
-address, which can be useful for troubleshooting network issues.
+This feature displays information about the public IP address used by
+devices to communicate with OpenWISP (via the ``last_ip`` field). It helps
+identify the geographic location and ISP associated with the IP address,
+which can be useful for troubleshooting network issues.
 
 The retrieved information pertains to the Autonomous System (ASN)
 associated with the device's public IP address and includes:
@@ -126,6 +126,5 @@ If a record is older than :ref:`Threshold
 <openwisp_controller_whois_refresh_threshold_days>`, it will be refreshed
 automatically.
 
-This will be triggered for the same scenarios defined in `trigger
-conditions <whois_trigger_conditions_>`_ but among the conditions **only
-WHOIS enablement is required**.
+The update mechanism will be triggered whenever a device is registered or
+its last IP changes or fetches its checksum.

openwisp_controller/config/management/commands/clear_last_ip.py

@@ -4,7 +4,10 @@
 
 
 class Command(BaseCommand):
-    help = "Clear the last IP address, if set, of active devices of all organizations."
+    help = (
+        "Clears the last IP address (if set) for every active device"
+        " across all organizations."
+    )
 
     def add_arguments(self, parser):
         parser.add_argument(
@@ -29,11 +32,11 @@ def handle(self, *args, **options):
                 "Are you sure you want to do this?\n\n"
                 "Type 'yes' to continue, or 'no' to cancel: "
             )
-            if input("".join(message)) != "yes":
+            if input("".join(message)).lower() != "yes":
                 raise CommandError("Operation cancelled by user.")
 
         devices = Device.objects.filter(_is_deactivated=False).only("last_ip")
-        # Filter devices that have no WHOIS information for their last IP
+        # Filter out devices that have WHOIS information for their last IP
         devices = devices.exclude(last_ip=None).exclude(
             last_ip__in=Subquery(
                 WHOISInfo.objects.filter(ip_address=OuterRef("last_ip")).values(

openwisp_controller/config/whois/service.py

@@ -1,15 +1,37 @@
 from datetime import timedelta
 from ipaddress import ip_address as ip_addr
 
+import requests
+from django.contrib.gis.geos import Point
 from django.core.cache import cache
 from django.db import transaction
 from django.utils import timezone
+from django.utils.translation import gettext as _
+from geoip2 import errors
+from geoip2 import webservice as geoip2_webservice
 from swapper import load_model
 
 from openwisp_controller.config import settings as app_settings
 
 from .tasks import fetch_whois_details, manage_estimated_locations
 
+EXCEPTION_MESSAGES = {
+    errors.AddressNotFoundError: _(
+        "No WHOIS information found for IP address {ip_address}"
+    ),
+    errors.AuthenticationError: _(
+        "Authentication failed for GeoIP2 service. "
+        "Check your OPENWISP_CONTROLLER_WHOIS_GEOIP_ACCOUNT and "
+        "OPENWISP_CONTROLLER_WHOIS_GEOIP_KEY settings."
+    ),
+    errors.OutOfQueriesError: _(
+        "Your account has run out of queries for the GeoIP2 service."
+    ),
+    errors.PermissionRequiredError: _(
+        "Your account does not have permission to access this service."
+    ),
+}
+
 
 class WHOISService:
     """
@@ -19,6 +41,20 @@ class WHOISService:
     def __init__(self, device):
         self.device = device
 
+    @staticmethod
+    def geoip_client():
+        """
+        Used to get a GeoIP2 web service client instance.
+        Host is based on the db that is used to fetch the details.
+        As we are using GeoLite2, 'geolite.info' host is used.
+        Refer: https://geoip2.readthedocs.io/en/latest/#sync-web-service-example
+        """
+        return geoip2_webservice.Client(
+            account_id=app_settings.WHOIS_GEOIP_ACCOUNT,
+            license_key=app_settings.WHOIS_GEOIP_KEY,
+            host="geolite.info",
+        )
+
     @staticmethod
     def get_cache_key(org_id):
         """
@@ -116,6 +152,55 @@ def is_estimated_location_enabled(self):
         org_settings = self.get_org_config_settings(org_id=self.device.organization.pk)
         return org_settings.estimated_location_enabled
 
+    def process_whois_details(self, ip_address):
+        """
+        Fetch WHOIS details for a given IP address and return only
+        the relevant information.
+        """
+        ip_client = self.geoip_client()
+
+        try:
+            data = ip_client.city(ip_address=ip_address)
+
+        # Catching all possible exceptions raised by the geoip2 client
+        # and raising them with appropriate messages to be handled by the task
+        # retry mechanism.
+        except (
+            errors.AddressNotFoundError,
+            errors.AuthenticationError,
+            errors.OutOfQueriesError,
+            errors.PermissionRequiredError,
+        ) as e:
+            exc_type = type(e)
+            message = EXCEPTION_MESSAGES.get(exc_type)
+            if exc_type is errors.AddressNotFoundError:
+                message = message.format(ip_address=ip_address)
+            raise exc_type(message)
+        except requests.RequestException as e:
+            raise e
+
+        else:
+            # The attributes are always present in the response,
+            # but they can be None, so added fallbacks.
+            address = {
+                "city": data.city.name or "",
+                "country": data.country.name or "",
+                "continent": data.continent.name or "",
+                "postal": str(data.postal.code or ""),
+            }
+            coordinates = Point(
+                data.location.longitude, data.location.latitude, srid=4326
+            )
+            return {
+                "isp": data.traits.autonomous_system_organization,
+                "asn": data.traits.autonomous_system_number,
+                "timezone": data.location.time_zone,
+                "address": address,
+                "coordinates": coordinates,
+                "cidr": data.traits.network,
+                "ip_address": ip_address,
+            }
+
     def _need_whois_lookup(self, new_ip):
         """
         This is used to determine if the WHOIS lookup should be triggered

openwisp_controller/config/whois/tasks.py

@@ -1,12 +1,8 @@
 import logging
 
-import requests
 from celery import shared_task
-from django.contrib.gis.geos import Point
 from django.db import transaction
-from django.utils.translation import gettext as _
 from geoip2 import errors
-from geoip2 import webservice as geoip2_webservice
 from swapper import load_model
 
 from openwisp_controller.geo.estimated_location.tasks import manage_estimated_locations
@@ -17,23 +13,6 @@
 
 logger = logging.getLogger(__name__)
 
-EXCEPTION_MESSAGES = {
-    errors.AddressNotFoundError: _(
-        "No WHOIS information found for IP address {ip_address}"
-    ),
-    errors.AuthenticationError: _(
-        "Authentication failed for GeoIP2 service. "
-        "Check your OPENWISP_CONTROLLER_WHOIS_GEOIP_ACCOUNT and "
-        "OPENWISP_CONTROLLER_WHOIS_GEOIP_KEY settings."
-    ),
-    errors.OutOfQueriesError: _(
-        "Your account has run out of queries for the GeoIP2 service."
-    ),
-    errors.PermissionRequiredError: _(
-        "Your account does not have permission to access this service."
-    ),
-}
-
 
 class WHOISCeleryRetryTask(OpenwispCeleryTask):
     """
@@ -54,6 +33,28 @@ def on_failure(self, exc, task_id, args, kwargs, einfo):
         return super().on_failure(exc, task_id, args, kwargs, einfo)
 
 
+def _manage_whois_record(whois_details, whois_instance=None):
+    """
+    Used to update an existing WHOIS instance; else, creates a new one.
+    Returns the updated or created WHOIS instance along with update fields.
+    """
+    WHOISInfo = load_model("config", "WHOISInfo")
+
+    update_fields = []
+    if whois_instance:
+        for attr, value in whois_details.items():
+            if getattr(whois_instance, attr) != value:
+                update_fields.append(attr)
+                setattr(whois_instance, attr, value)
+        if update_fields:
+            whois_instance.save(update_fields=update_fields)
+    else:
+        whois_instance = WHOISInfo(**whois_details)
+        whois_instance.full_clean()
+        whois_instance.save()
+    return whois_instance, update_fields
+
+
 # device_pk is used when task fails to report for which device failure occurred
 @shared_task(
     bind=True,
@@ -71,93 +72,35 @@ def fetch_whois_details(self, device_pk, initial_ip_address):
     with transaction.atomic():
         device = Device.objects.get(pk=device_pk)
         new_ip_address = device.last_ip
+        WHOISService = device.whois_service
+
         # If there is existing WHOIS older record then it needs to be updated
         whois_obj = WHOISInfo.objects.filter(ip_address=new_ip_address).first()
-        if whois_obj and not device.whois_service.is_older(whois_obj.modified):
+        if whois_obj and not WHOISService.is_older(whois_obj.modified):
             return
 
-        # Host is based on the db that is used to fetch the details.
-        # As we are using GeoLite2, 'geolite.info' host is used.
-        # Refer: https://geoip2.readthedocs.io/en/latest/#sync-web-service-example
-        ip_client = geoip2_webservice.Client(
-            account_id=app_settings.WHOIS_GEOIP_ACCOUNT,
-            license_key=app_settings.WHOIS_GEOIP_KEY,
-            host="geolite.info",
-        )
+        fetched_details = WHOISService.process_whois_details(new_ip_address)
+        whois_obj, update_fields = _manage_whois_record(fetched_details, whois_obj)
+        logger.info(f"Successfully fetched WHOIS details for {new_ip_address}.")
 
-        try:
-            data = ip_client.city(ip_address=new_ip_address)
-
-        # Catching all possible exceptions raised by the geoip2 client
-        # and raising them with appropriate messages to be handled by the task
-        # retry mechanism.
-        except (
-            errors.AddressNotFoundError,
-            errors.AuthenticationError,
-            errors.OutOfQueriesError,
-            errors.PermissionRequiredError,
-        ) as e:
-            exc_type = type(e)
-            message = EXCEPTION_MESSAGES.get(exc_type)
-            if exc_type is errors.AddressNotFoundError:
-                message = message.format(ip_address=new_ip_address)
-            raise exc_type(message)
-        except requests.RequestException as e:
-            raise e
-
-        else:
-            # The attributes are always present in the response,
-            # but they can be None, so added fallbacks.
-            address = {
-                "city": data.city.name or "",
-                "country": data.country.name or "",
-                "continent": data.continent.name or "",
-                "postal": str(data.postal.code or ""),
-            }
-            coordinates = Point(
-                data.location.longitude, data.location.latitude, srid=4326
-            )
-            details = {
-                "isp": data.traits.autonomous_system_organization,
-                "asn": data.traits.autonomous_system_number,
-                "timezone": data.location.time_zone,
-                "address": address,
-                "coordinates": coordinates,
-                "cidr": data.traits.network,
-                "ip_address": new_ip_address,
-            }
-            update_fields = []
-            if whois_obj:
-                for attr, value in details.items():
-                    if getattr(whois_obj, attr) != value:
-                        update_fields.append(attr)
-                        setattr(whois_obj, attr, value)
-                if update_fields:
-                    whois_obj.save(update_fields=update_fields)
-            else:
-                whois_obj = WHOISInfo(**details)
-                whois_obj.full_clean()
-                whois_obj.save()
-            logger.info(f"Successfully fetched WHOIS details for {new_ip_address}.")
-
-            if device._get_organization__config_settings().estimated_location_enabled:
-                # the estimated location task should not run if old record is updated
-                # and location related fields are not updated
-                if update_fields and not any(
-                    i in update_fields for i in ["address", "coordinates"]
-                ):
-                    return
-                manage_estimated_locations.delay(
-                    device_pk=device_pk, ip_address=new_ip_address
-                )
-
-            # delete WHOIS record for initial IP if no devices are linked to it
-            if (
-                not Device.objects.filter(_is_deactivated=False)
-                .filter(last_ip=initial_ip_address)
-                .exists()
+        if device._get_organization__config_settings().estimated_location_enabled:
+            # the estimated location task should not run if old record is updated
+            # and location related fields are not updated
+            if update_fields and not any(
+                i in update_fields for i in ["address", "coordinates"]
             ):
-                delete_whois_record(ip_address=initial_ip_address)
+                return
+            manage_estimated_locations.delay(
+                device_pk=device_pk, ip_address=new_ip_address
+            )
+
+        # delete WHOIS record for initial IP if no devices are linked to it
+        if (
+            not Device.objects.filter(_is_deactivated=False)
+            .filter(last_ip=initial_ip_address)
+            .exists()
+        ):
+            delete_whois_record(ip_address=initial_ip_address)
 
 
 @shared_task

openwisp_controller/config/whois/tests/tests.py

@@ -343,7 +343,7 @@ class TestWHOISTransaction(
     CreateWHOISMixin, WHOISTransactionMixin, TransactionTestCase
 ):
     _WHOIS_GEOIP_CLIENT = (
-        "openwisp_controller.config.whois.tasks.geoip2_webservice.Client"
+        "openwisp_controller.config.whois.service.geoip2_webservice.Client"
     )
     _WHOIS_TASKS_INFO_LOGGER = "openwisp_controller.config.whois.tasks.logger.info"
     _WHOIS_TASKS_WARN_LOGGER = "openwisp_controller.config.whois.tasks.logger.warning"