policies/reputation: save to database directly (#10059)

* policies/reputation: save to database directly

Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space>

* makemigrations

Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space>

* fix settings

Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space>

* also update expiry

Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space>

* lint?

Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space>

---------

Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space>

authentik/lib/default.yml

@@ -50,7 +50,6 @@ cache:
   timeout: 300
   timeout_flows: 300
   timeout_policies: 300
-  timeout_reputation: 300
 
 # channel:
 #   url: ""

authentik/policies/reputation/apps.py

@@ -2,8 +2,6 @@
 
 from authentik.blueprints.apps import ManagedAppConfig
 
-CACHE_KEY_PREFIX = "goauthentik.io/policies/reputation/scores/"
-
 
 class AuthentikPolicyReputationConfig(ManagedAppConfig):
     """Authentik reputation app config"""

authentik/policies/reputation/migrations/0007_reputation_authentik_p_identif_9434d7_idx_and_more.py

@@ -0,0 +1,25 @@
+# Generated by Django 5.0.6 on 2024-06-11 08:50
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("authentik_policies_reputation", "0006_reputation_ip_asn_data"),
+    ]
+
+    operations = [
+        migrations.AddIndex(
+            model_name="reputation",
+            index=models.Index(fields=["identifier"], name="authentik_p_identif_9434d7_idx"),
+        ),
+        migrations.AddIndex(
+            model_name="reputation",
+            index=models.Index(fields=["ip"], name="authentik_p_ip_7ad0df_idx"),
+        ),
+        migrations.AddIndex(
+            model_name="reputation",
+            index=models.Index(fields=["ip", "identifier"], name="authentik_p_ip_d779aa_idx"),
+        ),
+    ]

authentik/policies/reputation/models.py

@@ -96,3 +96,8 @@ class Reputation(ExpiringModel, SerializerModel):
         verbose_name = _("Reputation Score")
         verbose_name_plural = _("Reputation Scores")
         unique_together = ("identifier", "ip")
+        indexes = [
+            models.Index(fields=["identifier"]),
+            models.Index(fields=["ip"]),
+            models.Index(fields=["ip", "identifier"]),
+        ]

authentik/policies/reputation/settings.py

@@ -1,11 +0,0 @@
-"""Reputation Settings"""
-
-from celery.schedules import crontab
-
-CELERY_BEAT_SCHEDULE = {
-    "policies_reputation_save": {
-        "task": "authentik.policies.reputation.tasks.save_reputation",
-        "schedule": crontab(minute="1-59/5"),
-        "options": {"queue": "authentik_scheduled"},
-    },
-}

authentik/policies/reputation/signals.py

@@ -1,40 +1,35 @@
 """authentik reputation request signals"""
 
 from django.contrib.auth.signals import user_logged_in
-from django.core.cache import cache
 from django.dispatch import receiver
 from django.http import HttpRequest
 from structlog.stdlib import get_logger
 
 from authentik.core.signals import login_failed
-from authentik.lib.config import CONFIG
+from authentik.events.context_processors.asn import ASN_CONTEXT_PROCESSOR
-from authentik.policies.reputation.apps import CACHE_KEY_PREFIX
+from authentik.events.context_processors.geoip import GEOIP_CONTEXT_PROCESSOR
-from authentik.policies.reputation.tasks import save_reputation
+from authentik.policies.reputation.models import Reputation, reputation_expiry
 from authentik.root.middleware import ClientIPMiddleware
 from authentik.stages.identification.signals import identification_failed
 
 LOGGER = get_logger()
-CACHE_TIMEOUT = CONFIG.get_int("cache.timeout_reputation")
 
 
 def update_score(request: HttpRequest, identifier: str, amount: int):
     """Update score for IP and User"""
     remote_ip = ClientIPMiddleware.get_client_ip(request)
 
-    try:
+    Reputation.objects.update_or_create(
-        # We only update the cache here, as its faster than writing to the DB
+        ip=remote_ip,
-        score = cache.get_or_set(
+        identifier=identifier,
-            CACHE_KEY_PREFIX + remote_ip + "/" + identifier,
+        defaults={
-            {"ip": remote_ip, "identifier": identifier, "score": 0},
+            "score": amount,
-            CACHE_TIMEOUT,
+            "ip_geo_data": GEOIP_CONTEXT_PROCESSOR.city_dict(remote_ip) or {},
-        )
+            "ip_asn_data": ASN_CONTEXT_PROCESSOR.asn_dict(remote_ip) or {},
-        score["score"] += amount
+            "expires": reputation_expiry(),
-        cache.set(CACHE_KEY_PREFIX + remote_ip + "/" + identifier, score)
+        },
-    except ValueError as exc:
+    )
-        LOGGER.warning("failed to set reputation", exc=exc)
-
     LOGGER.debug("Updated score", amount=amount, for_user=identifier, for_ip=remote_ip)
-    save_reputation.delay()
 
 
 @receiver(login_failed)

authentik/policies/reputation/tasks.py

@@ -1,32 +0,0 @@
-"""Reputation tasks"""
-
-from django.core.cache import cache
-from structlog.stdlib import get_logger
-
-from authentik.events.context_processors.asn import ASN_CONTEXT_PROCESSOR
-from authentik.events.context_processors.geoip import GEOIP_CONTEXT_PROCESSOR
-from authentik.events.models import TaskStatus
-from authentik.events.system_tasks import SystemTask, prefill_task
-from authentik.policies.reputation.apps import CACHE_KEY_PREFIX
-from authentik.policies.reputation.models import Reputation
-from authentik.root.celery import CELERY_APP
-
-LOGGER = get_logger()
-
-
-@CELERY_APP.task(bind=True, base=SystemTask)
-@prefill_task
-def save_reputation(self: SystemTask):
-    """Save currently cached reputation to database"""
-    objects_to_update = []
-    for _, score in cache.get_many(cache.keys(CACHE_KEY_PREFIX + "*")).items():
-        rep, _ = Reputation.objects.get_or_create(
-            ip=score["ip"],
-            identifier=score["identifier"],
-        )
-        rep.ip_geo_data = GEOIP_CONTEXT_PROCESSOR.city_dict(score["ip"]) or {}
-        rep.ip_asn_data = ASN_CONTEXT_PROCESSOR.asn_dict(score["ip"]) or {}
-        rep.score = score["score"]
-        objects_to_update.append(rep)
-    Reputation.objects.bulk_update(objects_to_update, ["score", "ip_geo_data"])
-    self.set_status(TaskStatus.SUCCESSFUL, "Successfully updated Reputation")

authentik/policies/reputation/tests.py

@@ -1,14 +1,11 @@
 """test reputation signals and policy"""
 
-from django.core.cache import cache
 from django.test import RequestFactory, TestCase
 
 from authentik.core.models import User
 from authentik.lib.generators import generate_id
 from authentik.policies.reputation.api import ReputationPolicySerializer
-from authentik.policies.reputation.apps import CACHE_KEY_PREFIX
 from authentik.policies.reputation.models import Reputation, ReputationPolicy
-from authentik.policies.reputation.tasks import save_reputation
 from authentik.policies.types import PolicyRequest
 from authentik.stages.password import BACKEND_INBUILT
 from authentik.stages.password.stage import authenticate
@@ -22,8 +19,6 @@ class TestReputationPolicy(TestCase):
         self.request = self.request_factory.get("/")
         self.test_ip = "127.0.0.1"
         self.test_username = "test"
-        keys = cache.keys(CACHE_KEY_PREFIX + "*")
-        cache.delete_many(keys)
         # We need a user for the one-to-one in userreputation
         self.user = User.objects.create(username=self.test_username)
         self.backends = [BACKEND_INBUILT]
@@ -34,13 +29,6 @@ class TestReputationPolicy(TestCase):
         authenticate(
             self.request, self.backends, username=self.test_username, password=self.test_username
         )
-        # Test value in cache
-        self.assertEqual(
-            cache.get(CACHE_KEY_PREFIX + self.test_ip + "/" + self.test_username),
-            {"ip": "127.0.0.1", "identifier": "test", "score": -1},
-        )
-        # Save cache and check db values
-        save_reputation.delay().get()
         self.assertEqual(Reputation.objects.get(ip=self.test_ip).score, -1)
 
     def test_user_reputation(self):
@@ -49,13 +37,6 @@ class TestReputationPolicy(TestCase):
         authenticate(
             self.request, self.backends, username=self.test_username, password=self.test_username
         )
-        # Test value in cache
-        self.assertEqual(
-            cache.get(CACHE_KEY_PREFIX + self.test_ip + "/" + self.test_username),
-            {"ip": "127.0.0.1", "identifier": "test", "score": -1},
-        )
-        # Save cache and check db values
-        save_reputation.delay().get()
         self.assertEqual(Reputation.objects.get(identifier=self.test_username).score, -1)
 
     def test_policy(self):