diff --git a/passbook/policies/reputation/models.py b/passbook/policies/reputation/models.py index a4b22dff9..aa4326129 100644 --- a/passbook/policies/reputation/models.py +++ b/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: diff --git a/passbook/policies/reputation/settings.py b/passbook/policies/reputation/settings.py new file mode 100644 index 000000000..cc848ec90 --- /dev/null +++ b/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"), + }, +} diff --git a/passbook/policies/reputation/signals.py b/passbook/policies/reputation/signals.py index e11c38a93..e7daa8141 100644 --- a/passbook/policies/reputation/signals.py +++ b/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) diff --git a/passbook/policies/reputation/tasks.py b/passbook/policies/reputation/tasks.py new file mode 100644 index 000000000..916278d94 --- /dev/null +++ b/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"]) diff --git a/passbook/policies/reputation/tests.py b/passbook/policies/reputation/tests.py new file mode 100644 index 000000000..20b643bef --- /dev/null +++ b/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)