[enhancement] Added task and logic for fuzzy location creation #1034

Fuzzy location task should be executed when
creating WhoIs for a device. It should also
be executed when we have an existing WhoIs
and attach the device to the same location
as the similar ip's devices are attached.

Closes #1034

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

Author: DragnEmperor

openwisp_controller/config/whois/service.py

@@ -6,7 +6,7 @@
 
 from openwisp_controller.config import settings as app_settings
 
-from .tasks import fetch_whois_details
+from .tasks import fetch_whois_details, manage_fuzzy_locations
 
 
 class WHOISService:
@@ -79,18 +79,23 @@ def _need_whois_lookup(self, new_ip):
 
         The lookup is not triggered if:
             - The new IP address is None or it is a private IP address.
-            - The WHOIS information of new ip is already present.
-            - WHOIS is disabled in the organization settings. (query from db)
+            - The WhoIs information of new ip is already present.
+            - WhoIs is disabled in the organization settings. (query from db)
+
+        Two boolean values are returned:
+            - First boolean indicates if WhoIs lookup is needed.
+            - Second boolean indicates if WhoIs info already exists in the db,
+            which is used for managing fuzzy locations.
         """
 
         # Check cheap conditions first before hitting the database
         if not self.is_valid_public_ip_address(new_ip):
-            return False
+            return False, False
 
         if self._get_whois_info_from_db(new_ip).exists():
-            return False
+            return False, True
 
-        return self.is_whois_enabled
+        return self.is_whois_enabled, False
 
     def get_device_whois_info(self):
         """
@@ -108,11 +113,19 @@ def trigger_whois_lookup(self):
         Trigger WHOIS lookup based on the conditions of `_need_whois_lookup`.
         Task is triggered on commit to ensure redundant data is not created.
         """
-        if self._need_whois_lookup(self.device.last_ip):
+
+        fetch_whois, whois_info_exists = self._need_whois_lookup(self.device.last_ip)
+        if fetch_whois:
             transaction.on_commit(
                 lambda: fetch_whois_details.delay(
                     device_pk=self.device.pk,
                     initial_ip_address=self.device._initial_last_ip,
                     new_ip_address=self.device.last_ip,
                 )
             )
+        elif whois_info_exists and self.is_whois_enabled:
+            transaction.on_commit(
+                manage_fuzzy_locations.delay(
+                    self.device.pk, self.device.last_ip, add_existing=True
+                )
+            )

openwisp_controller/config/whois/tasks.py

@@ -2,6 +2,7 @@
 
 import requests
 from celery import shared_task
+from django.contrib.gis.geos import Point
 from django.utils.translation import gettext as _
 from geoip2 import errors
 from geoip2 import webservice as geoip2_webservice
@@ -158,3 +159,72 @@ def delete_whois_record(ip_address):
     queryset = WHOISInfo.objects.filter(ip_address=ip_address)
     if queryset.exists():
         queryset.delete()
+
+
+@shared_task
+def manage_fuzzy_locations(
+    device_pk,
+    ip_address,
+    latitude=None,
+    longitude=None,
+    address=None,
+    add_existing=False,
+):
+    """
+    Creates/updates fuzzy location for a device based on the latitude and longitude
+    or attaches an existing location if `add_existing` is True.
+    Existing location here means a location of a device whose last_ip matches
+    the given ip_address.
+    """
+    Device = load_model("config", "Device")
+    Location = load_model("geo", "Location")
+    DeviceLocation = load_model("geo", "DeviceLocation")
+
+    device_location = (
+        DeviceLocation.objects.filter(content_object_id=device_pk)
+        .select_related("location")
+        .first()
+    )
+
+    if not device_location:
+        device_location = DeviceLocation(content_object_id=device_pk)
+
+    # if attaching an existing location, the current device should not have
+    # a location already set.
+    # TODO: Do we do this if device location exists but is marked fuzzy?
+    if add_existing and not device_location.location:
+        existing_device_with_location = (
+            Device.objects.select_related("device_location")
+            .filter(last_ip=ip_address, device_location__location__isnull=False)
+            .first()
+        )
+        if existing_device_with_location:
+            location = existing_device_with_location.device_location.location
+            device_location.location = location
+            device_location.full_clean()
+            device_location.save()
+    elif latitude and longitude:
+        device = Device.objects.get(pk=device_pk)
+        coords = Point(longitude, latitude, srid=4326)
+        # Create/update the device location mapping, updating existing location
+        # if exists else create a new location
+        location_defaults = {
+            "name": f"{device.name} Location",
+            "type": "outdoor",
+            "organization_id": device.organization_id,
+            "is_mobile": False,
+            "geometry": coords,
+            "address": address,
+        }
+        if device_location.location and device_location.location.fuzzy:
+            for attr, value in location_defaults.items():
+                setattr(device_location.location, attr, value)
+            device_location.location.full_clean()
+            device_location.location.save()
+        elif not device_location.location:
+            location = Location(**location_defaults, fuzzy=True)
+            location.full_clean()
+            location.save()
+            device_location.location = location
+            device_location.full_clean()
+            device_location.save()