From fcfa88d67905a9534a143f5ebeab35148593e534 Mon Sep 17 00:00:00 2001
From: horida <unknown>
Date: Fri, 6 Sep 2019 10:47:27 +0200
Subject: [PATCH] Add unblock signals

---
 defender/signals.py | 13 +++++++++++++
 defender/tests.py   | 24 +++++++++++++++++++++++-
 defender/utils.py   |  5 ++++-
 3 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/defender/signals.py b/defender/signals.py
index 4aaaec3..c9cd9c1 100644
--- a/defender/signals.py
+++ b/defender/signals.py
@@ -1,7 +1,10 @@
 from django.dispatch import Signal
 
 username_block = Signal(providing_args=['username'])
+username_unblock = Signal(providing_args=['username'])
 ip_block = Signal(providing_args=['ip_address'])
+ip_unblock = Signal(providing_args=['ip_address'])
+
 
 class BlockSignal:
     """
@@ -10,8 +13,18 @@ class BlockSignal:
     """
     pass
 
+
 def send_username_block_signal(username):
     username_block.send(sender=BlockSignal, username=username)
 
+
+def send_username_unblock_signal(username):
+    username_unblock.send(sender=BlockSignal, username=username)
+
+
 def send_ip_block_signal(ip_address):
     ip_block.send(sender=BlockSignal, ip_address=ip_address)
+
+
+def send_ip_unblock_signal(ip_address):
+    ip_unblock.send(sender=BlockSignal, ip_address=ip_address)
diff --git a/defender/tests.py b/defender/tests.py
index c2bc920..7b07d8a 100644
--- a/defender/tests.py
+++ b/defender/tests.py
@@ -25,7 +25,9 @@ from . import utils
 from . import config
 from .signals import (
     ip_block as ip_block_signal,
-    username_block as username_block_signal
+    ip_unblock as ip_unblock_signal,
+    username_block as username_block_signal,
+    username_unblock as username_unblock_signal,
 )
 from .connection import parse_redis_url, get_redis_connection
 from .decorators import watch_login
@@ -936,6 +938,16 @@ class SignalTest(DefenderTestCase):
         utils.block_ip('8.8.8.8')
         self.assertEqual(self.blocked_ip, '8.8.8.8')
 
+    def test_should_send_signal_when_unblocking_ip(self):
+        self.blocked_ip = ('8.8.8.8')
+
+        def handler(sender, ip_address, **kwargs):
+            self.blocked_ip = None
+
+        ip_unblock_signal.connect(handler)
+        utils.unblock_ip('8.8.8.8')
+        self.assertIsNone(self.blocked_ip)
+
     def test_should_send_signal_when_blocking_username(self):
         self.blocked_username = None
 
@@ -946,6 +958,16 @@ class SignalTest(DefenderTestCase):
         utils.block_username('richard_hendricks')
         self.assertEqual(self.blocked_username, 'richard_hendricks')
 
+    def test_should_send_signal_when_unblocking_username(self):
+        self.blocked_username = 'richard_hendricks'
+
+        def handler(sender, username, **kwargs):
+            self.blocked_username = None
+
+        username_unblock_signal.connect(handler)
+        utils.unblock_username('richard_hendricks')
+        self.assertIsNone(self.blocked_username)
+
 
 class DefenderTestCaseTest(DefenderTestCase):
     """ Make sure that we're cleaning the cache between tests """
diff --git a/defender/utils.py b/defender/utils.py
index 1952629..cb34cc4 100644
--- a/defender/utils.py
+++ b/defender/utils.py
@@ -10,7 +10,8 @@ from django.utils.module_loading import import_string
 from .connection import get_redis_connection
 from . import config
 from .data import store_login_attempt
-from .signals import send_username_block_signal, send_ip_block_signal
+from .signals import (send_username_block_signal, send_ip_block_signal,
+                      send_username_unblock_signal, send_ip_unblock_signal)
 
 REDIS_SERVER = get_redis_connection()
 
@@ -254,6 +255,7 @@ def unblock_ip(ip_address, pipe=None):
         pipe.delete(get_ip_blocked_cache_key(ip_address))
         if do_commit:
             pipe.execute()
+    send_ip_unblock_signal(ip_address)
 
 
 def unblock_username(username, pipe=None):
@@ -267,6 +269,7 @@ def unblock_username(username, pipe=None):
         pipe.delete(get_username_blocked_cache_key(username))
         if do_commit:
             pipe.execute()
+    send_username_unblock_signal(username)
 
 
 def reset_failed_attempts(ip_address=None, username=None):