From 40031bd29d012c477c88d74dd51a79db57add68c Mon Sep 17 00:00:00 2001
From: shaohuzhang1 <shaohu.zhang@fit2cloud.com>
Date: Wed, 20 Mar 2024 11:57:17 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E5=AF=86=E9=92=A5=E5=AF=B9=E5=AD=98?=
 =?UTF-8?q?=E5=82=A8=E5=88=B0=E6=95=B0=E6=8D=AE=E5=BA=93?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/common/util/rsa_util.py             | 50 ++++++++++++++++++++----
 apps/setting/models/system_management.py |  2 +
 2 files changed, 44 insertions(+), 8 deletions(-)

diff --git a/apps/common/util/rsa_util.py b/apps/common/util/rsa_util.py
index ed950776a..a8e361e8c 100644
--- a/apps/common/util/rsa_util.py
+++ b/apps/common/util/rsa_util.py
@@ -11,7 +11,16 @@ import os
 
 from Crypto.Cipher import PKCS1_v1_5 as PKCS1_cipher
 from Crypto.PublicKey import RSA
+from django.db.models import QuerySet
+from django.core import cache
 
+from setting.models import SystemSetting, SettingType
+
+import threading
+
+lock = threading.Lock()
+rsa_cache = cache.caches['default']
+cache_key = "rsa_key"
 # 对密钥加密的密码
 secret_code = "mac_kb_password"
 
@@ -31,15 +40,40 @@ def generate():
 
 
 def get_key_pair():
-    if not os.path.exists("/opt/maxkb/conf/receiver.pem"):
+    rsa_value = rsa_cache.get(cache_key)
+    if rsa_value is None:
+        lock.acquire()
+        rsa_value = rsa_cache.get(cache_key)
+        if rsa_value is not None:
+            return rsa_value
+        try:
+            rsa_value = get_key_pair_by_sql()
+            rsa_cache.set(cache_key, rsa_value)
+        finally:
+            lock.release()
+    return rsa_value
+
+
+def get_key_pair_by_sql():
+    system_setting = QuerySet(SystemSetting).filter(type=SettingType.RSA.value).first()
+    if system_setting is None:
         kv = generate()
-        private_file_out = open("/opt/maxkb/conf/private.pem", "wb")
-        private_file_out.write(kv.get('value'))
-        private_file_out.close()
-        receiver_file_out = open("/opt/maxkb/conf/receiver.pem", "wb")
-        receiver_file_out.write(kv.get('key'))
-        receiver_file_out.close()
-    return {'key': open("/opt/maxkb/conf/receiver.pem").read(), 'value': open("/opt/maxkb/conf/private.pem").read()}
+        system_setting = SystemSetting(type=SettingType.RSA.value,
+                                       meta={'key': kv.get('key').decode(), 'value': kv.get('value').decode()})
+        system_setting.save()
+    return system_setting.meta
+
+
+# def get_key_pair():
+#     if not os.path.exists("/opt/maxkb/conf/receiver.pem"):
+#         kv = generate()
+#         private_file_out = open("/opt/maxkb/conf/private.pem", "wb")
+#         private_file_out.write(kv.get('value'))
+#         private_file_out.close()
+#         receiver_file_out = open("/opt/maxkb/conf/receiver.pem", "wb")
+#         receiver_file_out.write(kv.get('key'))
+#         receiver_file_out.close()
+#     return {'key': open("/opt/maxkb/conf/receiver.pem").read(), 'value': open("/opt/maxkb/conf/private.pem").read()}
 
 
 def encrypt(msg, public_key: str | None = None):
diff --git a/apps/setting/models/system_management.py b/apps/setting/models/system_management.py
index a4f4273f0..84060b792 100644
--- a/apps/setting/models/system_management.py
+++ b/apps/setting/models/system_management.py
@@ -17,6 +17,8 @@ class SettingType(models.IntegerChoices):
     """系统设置类型"""
     EMAIL = 0, '邮箱'
 
+    RSA = 1, "私钥秘钥"
+
 
 class SystemSetting(AppModelMixin):
     """