Delete old data task (certego#16) (certego#17)

* Fixed updated field for users risk_score * Added clear_models_periodically task * Refactoring * Added test for clear_models_periodically task Co-authored-by: Federico <[email protected]>
0x0elliot · Apr 5, 2023 · f76526c · f76526c
1 parent 0c86658
commit f76526c
Show file tree

Hide file tree

Showing 8 changed files with 103 additions and 23 deletions.
diff --git a/buffalogs/buffalogs/settings/settings.py b/buffalogs/buffalogs/settings/settings.py
@@ -182,6 +182,9 @@
 # Certego settings
 CERTEGO_DISTANCE_KM_ACCEPTED = 100
 CERTEGO_VEL_TRAVEL_ACCEPTED = 300
+CERTEGO_USER_MAX_DAYS = 20
+CERTEGO_LOGIN_MAX_DAYS = 10
+CERTEGO_ALERT_MAX_DAYS = 10
 
 # Celery config
 CELERY_BROKER_URL = CERTEGO_RABBITMQ_URI

diff --git a/buffalogs/impossible_travel/admin.py b/buffalogs/impossible_travel/admin.py
@@ -5,7 +5,7 @@
 
 @admin.register(Login)
 class LoginAdmin(admin.ModelAdmin):
-    list_display = ("id", "get_username", "timestamp", "latitude", "longitude", "country", "user_agent")
+    list_display = ("id", "created", "updated", "get_username", "timestamp", "latitude", "longitude", "country", "user_agent")
     search_fields = ("id", "user__username", "user_agent")
 
     @admin.display(description="username")
@@ -21,7 +21,7 @@ class UserAdmin(admin.ModelAdmin):
 
 @admin.register(Alert)
 class AlertAdmin(admin.ModelAdmin):
-    list_display = ("id", "get_username", "created", "name", "description", "login_raw_data")
+    list_display = ("id", "created", "updated", "get_username", "name", "description", "login_raw_data")
     search_fields = ("user__username", "name")
 
     @admin.display(description="username")

diff --git a/buffalogs/impossible_travel/management/commands/impossible_travel.py b/buffalogs/impossible_travel/management/commands/impossible_travel.py
@@ -7,7 +7,7 @@
 from django.utils import timezone
 from elasticsearch_dsl import Search, connections
 from impossible_travel.models import TaskSettings, User
-from impossible_travel.tasks import process_logs, process_user
+from impossible_travel.tasks import clear_models_periodically, process_logs, process_user, update_risk_level
 
 
 class Command(BaseCommand):
@@ -36,9 +36,12 @@ def handle(self, *args, **options):
             s.aggs.bucket("login_user", "terms", field="user.name", size=10000)
             response = s.execute()
 
+            clear_models_periodically()
+
             for user in response.aggregations.login_user.buckets:
                 db_user = User.objects.get_or_create(username=user.key)
                 process_user(db_user[0], start_date, end_date)
+                update_risk_level(db_user[0])
 
             if end_date >= now:
                 break
diff --git a/buffalogs/impossible_travel/migrations/0002_alert_updated.py b/buffalogs/impossible_travel/migrations/0002_alert_updated.py
@@ -0,0 +1,18 @@
+# Generated by Django 4.1.4 on 2023-04-03 10:15
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("impossible_travel", "0001_initial"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="alert",
+            name="updated",
+            field=models.DateField(auto_now=True),
+        ),
+    ]
diff --git a/buffalogs/impossible_travel/models.py b/buffalogs/impossible_travel/models.py
@@ -43,6 +43,7 @@ class ruleNameEnum(models.TextChoices):
     user = models.ForeignKey(User, on_delete=models.CASCADE)
     login_raw_data = models.JSONField()
     created = models.DateTimeField(auto_now_add=True)
+    updated = models.DateField(auto_now=True)
     description = models.TextField()
 
 

diff --git a/buffalogs/impossible_travel/modules/impossible_travel.py b/buffalogs/impossible_travel/modules/impossible_travel.py
@@ -1,10 +1,9 @@
 import logging
-from datetime import datetime
 
 from django.conf import settings
 from django.utils import timezone
 from geopy.distance import geodesic
-from impossible_travel.models import Alert, Login, User
+from impossible_travel.models import Alert, Login
 
 
 class Impossible_Travel:

diff --git a/buffalogs/impossible_travel/tasks.py b/buffalogs/impossible_travel/tasks.py
@@ -5,7 +5,6 @@
 from django.conf import settings
 from django.core.exceptions import ObjectDoesNotExist
 from django.db import transaction
-from django.db.models import Count
 from django.utils import timezone
 from elasticsearch_dsl import Search, connections
 from impossible_travel.models import Alert, Login, TaskSettings, User
@@ -14,21 +13,35 @@
 logger = logging.getLogger()
 
 
+def clear_models_periodically():
+    """
+    Clear DB models
+    """
+    now = timezone.now()
+    delete_user_time = now - timedelta(days=settings.CERTEGO_USER_MAX_DAYS)
+    User.objects.filter(updated__lte=delete_user_time).delete()
+
+    delete_login_time = now - timedelta(days=settings.CERTEGO_LOGIN_MAX_DAYS)
+    Login.objects.filter(updated__lte=delete_login_time).delete()
+
+    delete_alert_time = now - timedelta(days=settings.CERTEGO_ALERT_MAX_DAYS)
+    Alert.objects.filter(updated__lte=delete_alert_time).delete()
+
+
 @shared_task(name="UpdateRiskLevelTask")
-def update_risk_level():
+def update_risk_level(new_user):
     with transaction.atomic():
-        for u in User.objects.annotate(Count("alert")):
-            alerts_num = u.alert__count
-            if alerts_num == 0:
-                u.risk_score = User.riskScoreEnum.NO_RISK
-            elif 1 <= alerts_num <= 2:
-                u.risk_score = User.riskScoreEnum.LOW
-            elif 3 <= alerts_num <= 4:
-                u.risk_score = User.riskScoreEnum.MEDIUM
-            else:
-                logger.info(f"{User.riskScoreEnum.HIGH} risk level for User: {u.username}")
-                u.risk_score = User.riskScoreEnum.HIGH
-            u.save()
+        alerts_num = Alert.objects.filter(user__username=new_user.username).count()
+        if alerts_num == 0:
+            new_user.risk_score = User.riskScoreEnum.NO_RISK
+        elif 1 <= alerts_num <= 2:
+            new_user.risk_score = User.riskScoreEnum.LOW
+        elif 3 <= alerts_num <= 4:
+            new_user.risk_score = User.riskScoreEnum.MEDIUM
+        else:
+            logger.info(f"{User.riskScoreEnum.HIGH} risk level for User: {new_user.username}")
+            new_user.risk_score = User.riskScoreEnum.HIGH
+        new_user.save()
 
 
 def set_alert(db_user, login_alert, alert_info):

diff --git a/buffalogs/impossible_travel/tests/test_tasks.py b/buffalogs/impossible_travel/tests/test_tasks.py
@@ -1,5 +1,6 @@
 import json
 import os
+from datetime import timedelta
 from unittest.mock import patch
 
 from django.test import TestCase
@@ -66,15 +67,15 @@ def test_set_alert(self):
 
     def test_update_risk_level_norisk(self):
         # 0 alert --> no risk
-        tasks.update_risk_level()
         db_user = User.objects.get(username="Lorena Goldoni")
+        tasks.update_risk_level(db_user)
         self.assertEqual("No risk", db_user.risk_score)
 
     def test_update_risk_level_low(self):
         # 1 alert --> Low risk
         db_user = User.objects.get(username="Lorena Goldoni")
         Alert.objects.create(user=db_user, name=Alert.ruleNameEnum.IMP_TRAVEL, login_raw_data="Test", description="Test_Description")
-        tasks.update_risk_level()
+        tasks.update_risk_level(db_user)
         db_user = User.objects.get(username="Lorena Goldoni")
         self.assertEqual("Low", db_user.risk_score)
 
@@ -88,7 +89,7 @@ def test_update_risk_level_medium(self):
                 Alert(user=db_user, name=Alert.ruleNameEnum.NEW_COUNTRY, login_raw_data="Test3", description="Test_Description3"),
             ]
         )
-        tasks.update_risk_level()
+        tasks.update_risk_level(db_user)
         db_user = User.objects.get(username="Lorena Goldoni")
         self.assertEqual("Medium", db_user.risk_score)
 
@@ -104,7 +105,7 @@ def test_update_risk_level_high(self):
                 Alert(user=db_user, name=Alert.ruleNameEnum.NEW_COUNTRY, login_raw_data="Test5", description="Test_Description5"),
             ]
         )
-        tasks.update_risk_level()
+        tasks.update_risk_level(db_user)
         db_user = User.objects.get(username="Lorena Goldoni")
         self.assertEqual("High", db_user.risk_score)
 
@@ -122,3 +123,45 @@ def test_process_user(self, mock_execute, mock_chedk_fields):
         db_user = User.objects.get(username="Lorena Goldoni")
         tasks.process_user(db_user, iso_start_date, iso_end_date)
         mock_chedk_fields.assert_called_once_with(db_user, data_results)
+
+    def test_clear_models_periodically(self):
+        user_obj = User.objects.create(username="Lorena")
+        Login.objects.create(user=user_obj, timestamp=timezone.now())
+        raw_data = {
+            "lat": 40.6079,
+            "lon": -74.4037,
+            "agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko)",
+            "country": "United States",
+            "timestamp": "2023-04-03T14:01:47.907Z",
+        }
+        Alert.objects.create(user=user_obj, login_raw_data=raw_data)
+        old_date = timezone.now() + timedelta(days=-100)
+        User.objects.filter(username="Lorena").update(updated=old_date)
+        Login.objects.filter(user__username="Lorena").update(updated=old_date)
+        Alert.objects.filter(user__username="Lorena").update(updated=old_date)
+        tasks.clear_models_periodically()
+        with self.assertRaises(User.DoesNotExist):
+            User.objects.get(username="Lorena")
+        with self.assertRaises(Login.DoesNotExist):
+            Login.objects.get(user__username="Lorena")
+        with self.assertRaises(Alert.DoesNotExist):
+            Alert.objects.get(user__username="Lorena")
+
+    def test_clear_models_periodically_alert_delete(self):
+        user_obj = User.objects.create(username="Lorena")
+        Login.objects.create(user=user_obj, timestamp=timezone.now())
+        raw_data = {
+            "lat": 40.6079,
+            "lon": -74.4037,
+            "agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko)",
+            "country": "United States",
+            "timestamp": "2023-04-03T14:01:47.907Z",
+        }
+        Alert.objects.create(user=user_obj, login_raw_data=raw_data)
+        old_date = timezone.now() + timedelta(days=-100)
+        Alert.objects.filter(user__username="Lorena").update(updated=old_date)
+        tasks.clear_models_periodically()
+        with self.assertRaises(Alert.DoesNotExist):
+            Alert.objects.get(user__username="Lorena")
+        self.assertTrue(User.objects.filter(username="Lorena").exists())
+        self.assertTrue(Login.objects.filter(user__username="Lorena").exists())