From f76526c447306e05791602b859d6bb0e1bb9ddea Mon Sep 17 00:00:00 2001 From: Lorena Goldoni <33703137+Lorygold@users.noreply.github.com> Date: Wed, 5 Apr 2023 10:29:15 +0200 Subject: [PATCH] Delete old data task (#16) (#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 --- buffalogs/buffalogs/settings/settings.py | 3 ++ buffalogs/impossible_travel/admin.py | 4 +- .../management/commands/impossible_travel.py | 5 +- .../migrations/0002_alert_updated.py | 18 +++++++ buffalogs/impossible_travel/models.py | 1 + .../modules/impossible_travel.py | 3 +- buffalogs/impossible_travel/tasks.py | 41 ++++++++++----- .../impossible_travel/tests/test_tasks.py | 51 +++++++++++++++++-- 8 files changed, 103 insertions(+), 23 deletions(-) create mode 100644 buffalogs/impossible_travel/migrations/0002_alert_updated.py diff --git a/buffalogs/buffalogs/settings/settings.py b/buffalogs/buffalogs/settings/settings.py index f4ca2a0..19b31a4 100644 --- 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 index ce57372..3bbe755 100644 --- 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 index 3b99b50..9df6185 100644 --- 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 new file mode 100644 index 0000000..38e42ce --- /dev/null +++ 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 index 5b1cf92..7d193e5 100644 --- 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 index dcc39bd..d17bb7c 100644 --- 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 index 1634db0..18dfe43 100644 --- 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 index 0249fdc..2d1077a 100644 --- 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())