51 lines
1.9 KiB
Python
51 lines
1.9 KiB
Python
"""Reputation tasks"""
|
|
from django.core.cache import cache
|
|
from structlog import get_logger
|
|
|
|
from passbook.core.models import User
|
|
from passbook.lib.tasks import MonitoredTask, TaskResult, TaskResultStatus
|
|
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(bind=True, base=MonitoredTask)
|
|
def save_ip_reputation(self: MonitoredTask):
|
|
"""Save currently cached reputation to database"""
|
|
objects_to_update = []
|
|
for key, score in cache.get_many(cache.keys(CACHE_KEY_IP_PREFIX + "*")).items():
|
|
remote_ip = key.replace(CACHE_KEY_IP_PREFIX, "")
|
|
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"])
|
|
self.set_status(
|
|
TaskResult(TaskResultStatus.SUCCESSFUL, ["Successfully updated IP Reputation"])
|
|
)
|
|
|
|
|
|
@CELERY_APP.task(bind=True, base=MonitoredTask)
|
|
def save_user_reputation(self: MonitoredTask):
|
|
"""Save currently cached reputation to database"""
|
|
objects_to_update = []
|
|
for key, score in cache.get_many(cache.keys(CACHE_KEY_USER_PREFIX + "*")).items():
|
|
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"])
|
|
self.set_status(
|
|
TaskResult(
|
|
TaskResultStatus.SUCCESSFUL, ["Successfully updated User Reputation"]
|
|
)
|
|
)
|