policies/reputation: rewrite to save score into cache and save into DB via worker

passbook/policies/reputation/models.py

@@ -1,4 +1,5 @@
 """passbook reputation request policy"""
+from django.core.cache import cache
 from django.db import models
 from django.utils.translation import gettext as _
 
@@ -7,6 +8,9 @@ from passbook.lib.utils.http import get_client_ip
 from passbook.policies.models import Policy
 from passbook.policies.types import PolicyRequest, PolicyResult
 
+CACHE_KEY_IP_PREFIX = "passbook_reputation_ip_"
+CACHE_KEY_USER_PREFIX = "passbook_reputation_user_"
+
 
 class ReputationPolicy(Policy):
     """Return true if request IP/target username's score is below a certain threshold"""
@@ -18,18 +22,14 @@ class ReputationPolicy(Policy):
     form = "passbook.policies.reputation.forms.ReputationPolicyForm"
 
     def passes(self, request: PolicyRequest) -> PolicyResult:
-        remote_ip = get_client_ip(request.http_request)
+        remote_ip = get_client_ip(request.http_request) or "255.255.255.255"
         passing = True
         if self.check_ip:
-            ip_scores = IPReputation.objects.filter(
-                ip=remote_ip, score__lte=self.threshold
-            )
-            passing = passing and ip_scores.exists()
+            score = cache.get_or_set(CACHE_KEY_IP_PREFIX + remote_ip, 0)
+            passing = passing and score <= self.threshold
         if self.check_username:
-            user_scores = UserReputation.objects.filter(
-                user=request.user, score__lte=self.threshold
-            )
-            passing = passing and user_scores.exists()
+            score = cache.get_or_set(CACHE_KEY_USER_PREFIX + request.user.username, 0)
+            passing = passing and score <= self.threshold
         return PolicyResult(passing)
 
     class Meta:

passbook/policies/reputation/settings.py

@@ -0,0 +1,13 @@
+"""Reputation Settings"""
+from celery.schedules import crontab
+
+CELERY_BEAT_SCHEDULE = {
+    "policies_reputation_ip_save": {
+        "task": "passbook.policies.reputation.tasks.save_ip_reputation",
+        "schedule": crontab(minute="*/5"),
+    },
+    "policies_reputation_user_save": {
+        "task": "passbook.policies.reputation.tasks.save_user_reputation",
+        "schedule": crontab(minute="*/5"),
+    },
+}

passbook/policies/reputation/signals.py

@@ -1,29 +1,31 @@
 """passbook reputation request signals"""
 from django.contrib.auth.signals import user_logged_in, user_login_failed
+from django.core.cache import cache
 from django.dispatch import receiver
+from django.http import HttpRequest
 from structlog import get_logger
 
-from passbook.core.models import User
 from passbook.lib.utils.http import get_client_ip
-from passbook.policies.reputation.models import IPReputation, UserReputation
+from passbook.policies.reputation.models import (
+    CACHE_KEY_IP_PREFIX,
+    CACHE_KEY_USER_PREFIX,
+)
 
 LOGGER = get_logger()
 
 
-def update_score(request, username, amount):
+def update_score(request: HttpRequest, username: str, amount: int):
     """Update score for IP and User"""
-    remote_ip = get_client_ip(request) or "255.255.255.255."
-    ip_score, _ = IPReputation.objects.update_or_create(ip=remote_ip)
-    ip_score.score += amount
-    ip_score.save()
-    LOGGER.debug("Updated score", amount=amount, for_ip=remote_ip)
-    user = User.objects.filter(username=username)
-    if not user.exists():
-        return
-    user_score, _ = UserReputation.objects.update_or_create(user=user.first())
-    user_score.score += amount
-    user_score.save()
-    LOGGER.debug("Updated score", amount=amount, for_user=username)
+    remote_ip = get_client_ip(request) or "255.255.255.255"
+
+    # We only update the cache here, as its faster than writing to the DB
+    cache.get_or_set(CACHE_KEY_IP_PREFIX + remote_ip, 0)
+    cache.incr(CACHE_KEY_IP_PREFIX + remote_ip, amount)
+
+    cache.get_or_set(CACHE_KEY_USER_PREFIX + username, 0)
+    cache.incr(CACHE_KEY_USER_PREFIX + username, amount)
+
+    LOGGER.debug("Updated score", amount=amount, for_user=username, for_ip=remote_ip)
 
 
 @receiver(user_login_failed)

passbook/policies/reputation/tasks.py

@@ -0,0 +1,46 @@
+"""Reputation tasks"""
+from django.core.cache import cache
+from structlog import get_logger
+
+from passbook.core.models import User
+from passbook.policies.reputation.models import IPReputation, UserReputation
+from passbook.policies.reputation.signals import (
+    CACHE_KEY_IP_PREFIX,
+    CACHE_KEY_USER_PREFIX,
+)
+from passbook.root.celery import CELERY_APP
+
+LOGGER = get_logger()
+
+
+@CELERY_APP.task()
+def save_ip_reputation():
+    """Save currently cached reputation to database"""
+    keys = cache.keys(CACHE_KEY_IP_PREFIX + "*")
+    objects_to_update = []
+    for key in keys:
+        score = cache.get(key)
+        remote_ip = key.replace(CACHE_KEY_IP_PREFIX, "")
+        print(remote_ip)
+        rep, _ = IPReputation.objects.get_or_create(ip=remote_ip)
+        rep.score = score
+        objects_to_update.append(rep)
+    IPReputation.objects.bulk_update(objects_to_update, ["score"])
+
+
+@CELERY_APP.task()
+def save_user_reputation():
+    """Save currently cached reputation to database"""
+    keys = cache.keys(CACHE_KEY_USER_PREFIX + "*")
+    objects_to_update = []
+    for key in keys:
+        score = cache.get(key)
+        username = key.replace(CACHE_KEY_USER_PREFIX, "")
+        users = User.objects.filter(username=username)
+        if not users.exists():
+            LOGGER.info("User in cache does not exist, ignoring", username=username)
+            continue
+        rep, _ = UserReputation.objects.get_or_create(user=users.first())
+        rep.score = score
+        objects_to_update.append(rep)
+    UserReputation.objects.bulk_update(objects_to_update, ["score"])

passbook/policies/reputation/tests.py

@@ -0,0 +1,55 @@
+"""test reputation signals and policy"""
+from django.contrib.auth import authenticate
+from django.core.cache import cache
+from django.test import TestCase
+
+from passbook.core.models import User
+from passbook.policies.reputation.models import (
+    CACHE_KEY_IP_PREFIX,
+    CACHE_KEY_USER_PREFIX,
+    IPReputation,
+    ReputationPolicy,
+    UserReputation,
+)
+from passbook.policies.reputation.tasks import save_ip_reputation, save_user_reputation
+from passbook.policies.types import PolicyRequest
+
+
+class TestReputationPolicy(TestCase):
+    """test reputation signals and policy"""
+
+    def setUp(self):
+        self.test_ip = "255.255.255.255"
+        self.test_username = "test"
+        cache.delete(CACHE_KEY_IP_PREFIX + self.test_ip)
+        cache.delete(CACHE_KEY_USER_PREFIX + self.test_username)
+        # We need a user for the one-to-one in userreputation
+        self.user = User.objects.create(username=self.test_username)
+
+    def test_ip_reputation(self):
+        """test IP reputation"""
+        # Trigger negative reputation
+        authenticate(None, username=self.test_username, password=self.test_username)
+        # Test value in cache
+        self.assertEqual(cache.get(CACHE_KEY_IP_PREFIX + self.test_ip), -1)
+        # Save cache and check db values
+        save_ip_reputation()
+        self.assertEqual(IPReputation.objects.get(ip=self.test_ip).score, -1)
+
+    def test_user_reputation(self):
+        """test User reputation"""
+        # Trigger negative reputation
+        authenticate(None, username=self.test_username, password=self.test_username)
+        # Test value in cache
+        self.assertEqual(cache.get(CACHE_KEY_USER_PREFIX + self.test_username), -1)
+        # Save cache and check db values
+        save_user_reputation()
+        self.assertEqual(UserReputation.objects.get(user=self.user).score, -1)
+
+    def test_policy(self):
+        """Test Policy"""
+        request = PolicyRequest(user=self.user)
+        policy: ReputationPolicy = ReputationPolicy.objects.create(
+            name="reputation-test", threshold=0
+        )
+        self.assertTrue(policy.passes(request).passing)