注册码生成以及用户注册

elastic/es_connect.py

@@ -5,12 +5,13 @@ Django版本的ES连接和操作模块
 from elasticsearch import Elasticsearch
 from elasticsearch_dsl import connections
 import os
-from .documents import AchievementDocument, UserDocument, GlobalDocument
+from .documents import AchievementDocument, UserDocument, GlobalDocument, RegistrationCodeDocument
 from accounts.crypto import hash_password_random_salt
 from .indexes import ACHIEVEMENT_INDEX_NAME, USER_INDEX_NAME, GLOBAL_INDEX_NAME
 import hashlib
 import time
-from datetime import datetime, timezone
+from datetime import datetime, timezone, timedelta
+import uuid
 import json
 
 # 使用环境变量配置ES连接，默认为本机
@@ -115,6 +116,87 @@ def ensure_type_in_list(type_name: str):
     except Exception:
         return False
 
+def get_keys_list():
+    try:
+        try:
+            doc = GlobalDocument.get(id='keys')
+            cur = list(doc.keys_list or [])
+        except Exception:
+            cur = []
+            doc = GlobalDocument(keys_list=cur)
+            doc.meta.id = 'keys'
+            doc.save()
+        return [str(t).strip().strip(';') for t in cur]
+    except Exception:
+        return []
+
+def ensure_key_in_list(key_name: str):
+    if not key_name:
+        return False
+    norm = str(key_name).strip().strip(';')
+    try:
+        try:
+            doc = GlobalDocument.get(id='keys')
+            cur = list(doc.keys_list or [])
+        except Exception:
+            cur = []
+            doc = GlobalDocument(keys_list=cur)
+            doc.meta.id = 'keys'
+        cur_sanitized = {str(t).strip().strip(';') for t in cur}
+        if norm not in cur_sanitized:
+            cur.append(norm)
+            doc.keys_list = cur
+            doc.save()
+            return True
+        return False
+    except Exception:
+        return False
+
+def generate_registration_code(keys=None, manage_keys=None, expires_in_days: int = 30, created_by: int = None):
+    try:
+        keys = list(keys or [])
+        manage_keys = list(manage_keys or [])
+        for k in list(keys):
+            ensure_key_in_list(k)
+        for mk in list(manage_keys):
+            ensure_key_in_list(mk)
+        code = uuid.uuid4().hex + str(int(time.time()))[-6:]
+        now = datetime.now(timezone.utc)
+        expires = now + timedelta(days=max(1, int(expires_in_days or 30)))
+        doc = RegistrationCodeDocument(
+            code=code,
+            keys=keys,
+            manage_keys=manage_keys,
+            created_at=now.isoformat(),
+            expires_at=expires.isoformat(),
+            created_by=created_by,
+        )
+        doc.meta.id = code
+        doc.save()
+        return {
+            "code": code,
+            "keys": keys,
+            "manage_keys": manage_keys,
+            "created_at": now.isoformat(),
+            "expires_at": expires.isoformat(),
+        }
+    except Exception as e:
+        return None
+
+def get_registration_code(code: str):
+    try:
+        doc = RegistrationCodeDocument.get(id=str(code))
+        return {
+            "code": getattr(doc, 'code', str(code)),
+            "keys": list(getattr(doc, 'keys', []) or []),
+            "manage_keys": list(getattr(doc, 'manage_keys', []) or []),
+            "created_at": getattr(doc, 'created_at', None),
+            "expires_at": getattr(doc, 'expires_at', None),
+            "created_by": getattr(doc, 'created_by', None),
+        }
+    except Exception:
+        return None
+
 def get_doc_id(data):
     """
     根据数据内容生成唯一ID（用于去重）
@@ -527,7 +609,10 @@ def write_user_data(user_data):
             username=user_data.get('username'),
             password_hash=pwd_hash_b64,
             password_salt=pwd_salt_b64,
-            permission=perm_val
+            permission=perm_val,
+            email=user_data.get('email'),
+            key=list(user_data.get('key') or []),
+            manage_key=list(user_data.get('manage_key') or []),
         )
         user.save()
         print(f"用户数据写入成功: {user_data.get('username')}")