From 65f93951067f300d71b6b631d0caee82c43ec8d6 Mon Sep 17 00:00:00 2001
From: han0 <wslshanlin@sina.com>
Date: Thu, 19 Nov 2020 16:09:22 +0800
Subject: [PATCH] feat: verification

---
 nc_http/core/func.py                  | 15 +++++++
 nc_http/core/verification/__init__.py |  3 ++
 nc_http/core/verification/image.py    | 38 ++++++++++++++++
 nc_http/core/verification/storage.py  | 64 +++++++++++++++++++++++++++
 nc_http/tests/verification.py         | 17 +++++++
 5 files changed, 137 insertions(+)
 create mode 100644 nc_http/core/verification/__init__.py
 create mode 100644 nc_http/core/verification/image.py
 create mode 100644 nc_http/core/verification/storage.py
 create mode 100644 nc_http/tests/verification.py

diff --git a/nc_http/core/func.py b/nc_http/core/func.py
index 7153303..fb8dabf 100644
--- a/nc_http/core/func.py
+++ b/nc_http/core/func.py
@@ -73,3 +73,18 @@ def send_excel(file_handler, file_name, suffix='xlsx'):
         attachment_filename='{}.{}'.format(file_name, suffix),
         as_attachment=True
     )
+
+
+def send_png(img_handler, name='image'):
+    """
+    发送图片文件（可用于验证码）
+    :param img_handler:
+    :param name:
+    :return:
+    """
+    return send_file(
+        img_handler,
+        mimetype='image/png',
+        attachment_filename='{}.png'.format(name),
+        as_attachment=True
+    )
diff --git a/nc_http/core/verification/__init__.py b/nc_http/core/verification/__init__.py
new file mode 100644
index 0000000..1fd698f
--- /dev/null
+++ b/nc_http/core/verification/__init__.py
@@ -0,0 +1,3 @@
+"""
+用户验证码相关
+"""
\ No newline at end of file
diff --git a/nc_http/core/verification/image.py b/nc_http/core/verification/image.py
new file mode 100644
index 0000000..a4d569a
--- /dev/null
+++ b/nc_http/core/verification/image.py
@@ -0,0 +1,38 @@
+import random
+import string
+from io import BytesIO
+
+from captcha.image import ImageCaptcha
+
+
+class SimpleImageCaptcha:
+
+    def __init__(self, width=150, height=50, captcha_length=4, characters=string.digits, *args, **kwargs):
+        self.__width = width
+        self.__height = height
+        self.__captcha_length = captcha_length
+        self.__characters = characters
+
+        self.__str = ''.join([random.choice(characters) for _ in range(captcha_length)])
+
+        generator = ImageCaptcha(width=self.__width, height=self.__height, *args, **kwargs)
+        self.__image = generator.create_captcha_image(
+            chars=self.__str,
+            color=(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)),
+            background=(255, 255, 255)
+        )
+
+    @property
+    def str(self):
+        return self.__str
+
+    @property
+    def image(self):
+        return self.__image
+
+    @property
+    def image_handler(self):
+        output = BytesIO()
+        self.__image.save(output, 'png')
+        output.seek(0)
+        return output
diff --git a/nc_http/core/verification/storage.py b/nc_http/core/verification/storage.py
new file mode 100644
index 0000000..8bdd8d7
--- /dev/null
+++ b/nc_http/core/verification/storage.py
@@ -0,0 +1,64 @@
+from abc import abstractmethod, ABC
+
+from nc_http.core.verification.image import SimpleImageCaptcha
+
+
+class BaseCaptchaStorage(ABC):
+
+    @classmethod
+    def create(cls, *args, **kwargs):
+        """
+        创建验证码
+        :return:
+        """
+        captcha = SimpleImageCaptcha(*args, **kwargs)
+        cls.storage(captcha.str)
+        return captcha
+
+    @staticmethod
+    @abstractmethod
+    def storage(captcha_code):
+        """
+        记录生成的验证码字符
+        :param captcha_code:
+        :return:
+        """
+
+    @staticmethod
+    @abstractmethod
+    def validate(captcha_code):
+        """
+        确认验证码正确性
+        :param captcha_code:
+        :return: str
+        """
+
+
+class SampleCaptchaStorage(BaseCaptchaStorage):
+    """
+    验证码池样例实现
+
+        * 仅作测试使用，勿用于生产环境
+    """
+    storage_dict = {}
+
+    @staticmethod
+    def storage(captcha_code):
+        SampleCaptchaStorage.storage_dict[captcha_code] = True
+
+    @staticmethod
+    @abstractmethod
+    def validate(captcha_code):
+        valid = SampleCaptchaStorage.storage_dict.get(captcha_code)
+        if valid:
+            del SampleCaptchaStorage.storage_dict[captcha_code]
+            return True
+        return False
+
+
+if __name__ == '__main__':
+    test_captcha = SampleCaptchaStorage.create()
+    is_valid = SampleCaptchaStorage.validate('test_code')
+    print(is_valid)
+    is_valid = SampleCaptchaStorage.validate(test_captcha.str)
+    print(is_valid)
diff --git a/nc_http/tests/verification.py b/nc_http/tests/verification.py
new file mode 100644
index 0000000..3c60e34
--- /dev/null
+++ b/nc_http/tests/verification.py
@@ -0,0 +1,17 @@
+import unittest
+
+from nc_http.core.verification.storage import SampleCaptchaStorage
+
+
+class VerificationTestCase(unittest.TestCase):
+
+    def test_storage(self):
+        captcha = SampleCaptchaStorage.create()
+        is_valid = SampleCaptchaStorage.validate('test_code')
+        self.assertEqual(is_valid, False)
+        is_valid = SampleCaptchaStorage.validate(captcha.str)
+        self.assertEqual(is_valid, True)
+
+
+if __name__ == '__main__':
+    unittest.main()