diff --git a/defender/tests.py b/defender/tests.py index 037a4f9..0f71abe 100644 --- a/defender/tests.py +++ b/defender/tests.py @@ -1173,6 +1173,16 @@ class TestUtils(DefenderTestCase): req.META["HTTP_X_FORWARDED_FOR"] = "[2001:db8::1]:123456" self.assertEqual(utils.get_ip(req), "2001:db8::1") + def test_remove_prefix(self): + """ test the remove_prefix() method """ + self.assertEqual(utils.remove_prefix( + "defender:blocked:ip:192.168.24.24", "defender:blocked:"), "ip:192.168.24.24") + self.assertEqual(utils.remove_prefix( + "defender:blocked:username:johndoe", "defender:blocked:"), "username:johndoe") + self.assertEqual(utils.remove_prefix( + "defender:blocked:username:johndoe", "blocked:username:"), + "defender:blocked:username:johndoe") + class TestRedisConnection(TestCase): """ Test the redis connection parsing """ @@ -1224,4 +1234,4 @@ class TestRedisConnection(TestCase): redis_client.set('test3', 0) result = int(redis_client.get('test3')) self.assertEqual(result, 0) - redis_client.delete('test3') + redis_client.delete('test3') \ No newline at end of file diff --git a/defender/utils.py b/defender/utils.py index fb8ec3a..fb35fd0 100644 --- a/defender/utils.py +++ b/defender/utils.py @@ -1,6 +1,7 @@ from ipaddress import ip_address import logging import re +import sys from django.http import HttpResponse from django.http import HttpResponseRedirect @@ -127,20 +128,38 @@ def get_username_blocked_cache_key(username): ) +def remove_prefix(string, prefix): + if string.startswith(prefix): + return string[len(prefix):] + return string + + + def strip_keys(key_list): """ Given a list of keys, remove the prefix and remove just the data we care about. for example: - ['defender:blocked:ip:ken', 'defender:blocked:ip:joffrey'] + [ + 'defender:blocked:ip:192.168.24.24', + 'defender:blocked:ip:::ffff:192.168.24.24', + 'defender:blocked:username:joffrey' + ] would result in: - ['ken', 'joffrey'] - + [ + '192.168.24.24', + '::ffff:192.168.24.24', + 'joffrey' + ] """ - return [key.split(":")[-1] for key in key_list] + return [ + # key.removeprefix(f"{config.CACHE_PREFIX}:blocked:").partition(":")[2] + remove_prefix(key, f"{config.CACHE_PREFIX}:blocked:").partition(":")[2] + for key in key_list + ] def get_blocked_ips():