From 30e39c75ff9ddcdbb1eb95d3620c519338a51429 Mon Sep 17 00:00:00 2001 From: Marc 'risson' Schmitt Date: Fri, 14 Jun 2024 15:34:43 +0200 Subject: [PATCH] policies/reputation: save to database directly (#10059) * policies/reputation: save to database directly Signed-off-by: Marc 'risson' Schmitt * makemigrations Signed-off-by: Marc 'risson' Schmitt * fix settings Signed-off-by: Marc 'risson' Schmitt * also update expiry Signed-off-by: Marc 'risson' Schmitt * lint? Signed-off-by: Marc 'risson' Schmitt --------- Signed-off-by: Marc 'risson' Schmitt --- authentik/lib/default.yml | 1 - authentik/policies/reputation/apps.py | 2 -- ...authentik_p_identif_9434d7_idx_and_more.py | 25 +++++++++++++++ authentik/policies/reputation/models.py | 5 +++ authentik/policies/reputation/settings.py | 11 ------- authentik/policies/reputation/signals.py | 31 ++++++++---------- authentik/policies/reputation/tasks.py | 32 ------------------- authentik/policies/reputation/tests.py | 19 ----------- 8 files changed, 43 insertions(+), 83 deletions(-) create mode 100644 authentik/policies/reputation/migrations/0007_reputation_authentik_p_identif_9434d7_idx_and_more.py delete mode 100644 authentik/policies/reputation/settings.py delete mode 100644 authentik/policies/reputation/tasks.py diff --git a/authentik/lib/default.yml b/authentik/lib/default.yml index 5ca7a83756..cf6ca74f2c 100644 --- a/authentik/lib/default.yml +++ b/authentik/lib/default.yml @@ -50,7 +50,6 @@ cache: timeout: 300 timeout_flows: 300 timeout_policies: 300 - timeout_reputation: 300 # channel: # url: "" diff --git a/authentik/policies/reputation/apps.py b/authentik/policies/reputation/apps.py index 60b2c77b05..4a1b1509fc 100644 --- a/authentik/policies/reputation/apps.py +++ b/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""" diff --git a/authentik/policies/reputation/migrations/0007_reputation_authentik_p_identif_9434d7_idx_and_more.py b/authentik/policies/reputation/migrations/0007_reputation_authentik_p_identif_9434d7_idx_and_more.py new file mode 100644 index 0000000000..5395223e2b --- /dev/null +++ b/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"), + ), + ] diff --git a/authentik/policies/reputation/models.py b/authentik/policies/reputation/models.py index abcddac2cc..c3c7ab1ebf 100644 --- a/authentik/policies/reputation/models.py +++ b/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"]), + ] diff --git a/authentik/policies/reputation/settings.py b/authentik/policies/reputation/settings.py deleted file mode 100644 index 72abd27cf7..0000000000 --- a/authentik/policies/reputation/settings.py +++ /dev/null @@ -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"}, - }, -} diff --git a/authentik/policies/reputation/signals.py b/authentik/policies/reputation/signals.py index 5830143372..31042f8406 100644 --- a/authentik/policies/reputation/signals.py +++ b/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.policies.reputation.apps import CACHE_KEY_PREFIX -from authentik.policies.reputation.tasks import save_reputation +from authentik.events.context_processors.asn import ASN_CONTEXT_PROCESSOR +from authentik.events.context_processors.geoip import GEOIP_CONTEXT_PROCESSOR +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: - # We only update the cache here, as its faster than writing to the DB - score = cache.get_or_set( - CACHE_KEY_PREFIX + remote_ip + "/" + identifier, - {"ip": remote_ip, "identifier": identifier, "score": 0}, - CACHE_TIMEOUT, - ) - score["score"] += amount - cache.set(CACHE_KEY_PREFIX + remote_ip + "/" + identifier, score) - except ValueError as exc: - LOGGER.warning("failed to set reputation", exc=exc) - + Reputation.objects.update_or_create( + ip=remote_ip, + identifier=identifier, + defaults={ + "score": amount, + "ip_geo_data": GEOIP_CONTEXT_PROCESSOR.city_dict(remote_ip) or {}, + "ip_asn_data": ASN_CONTEXT_PROCESSOR.asn_dict(remote_ip) or {}, + "expires": reputation_expiry(), + }, + ) LOGGER.debug("Updated score", amount=amount, for_user=identifier, for_ip=remote_ip) - save_reputation.delay() @receiver(login_failed) diff --git a/authentik/policies/reputation/tasks.py b/authentik/policies/reputation/tasks.py deleted file mode 100644 index fcdceaa7e4..0000000000 --- a/authentik/policies/reputation/tasks.py +++ /dev/null @@ -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") diff --git a/authentik/policies/reputation/tests.py b/authentik/policies/reputation/tests.py index 534d808bf7..7d4e33fb79 100644 --- a/authentik/policies/reputation/tests.py +++ b/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):