尚未完成上传密钥时楼主身份验证,以及获取OTP时用户是否已购买附件的验证
import discord
from discord.ext import commands
import secrets
import string
import os
import requests
import re
import sqlite3
import pyotp
# 初始化 Bot
intents = discord.Intents.all()
bot = commands.Bot(command_prefix="!", intents=discord.Intents.all())
# SQLite数据库初始化
DATABASE_PATH = 'GM.db'
# 创建数据库并初始化表
def init_db():
with sqlite3.connect(DATABASE_PATH) as conn:
cursor = conn.cursor()
# 创建已验证用户表
cursor.execute('''
CREATE TABLE IF NOT EXISTS verified_users (
discord_user_id INTEGER NOT NULL,
uid TEXT NOT NULL,
PRIMARY KEY (discord_user_id)
)
''')
# 创建用户记录表
cursor.execute('''
CREATE TABLE IF NOT EXISTS user_records (
discord_user_id INTEGER NOT NULL,
tid TEXT NOT NULL,
zid TEXT NOT NULL,
secret TEXT NOT NULL,
PRIMARY KEY (discord_user_id, tid),
FOREIGN KEY (discord_user_id) REFERENCES verified_users (discord_user_id)
)
''')
conn.commit()
# 检查用户是否已验证,并返回对应的 uid
def is_user_verified(discord_user_id: int):
"""检查用户是否已验证,返回验证结果及 uid"""
with sqlite3.connect(DATABASE_PATH) as conn:
cursor = conn.cursor()
cursor.execute('SELECT uid FROM verified_users WHERE discord_user_id = ?', (discord_user_id,))
result = cursor.fetchone()
# 如果查询到结果,返回 True 和对应的 uid
if result:
return True, result[0]
else:
# 如果没有找到记录,返回 False 和 None
return False, None
# 检查 uid 是否已经被验证,并返回已绑定的 discord_user_id (如果有)
def is_uid_verified(uid: str):
"""检查指定的 uid 是否已经与其他 Discord 用户绑定"""
with sqlite3.connect(DATABASE_PATH) as conn:
cursor = conn.cursor()
cursor.execute('SELECT discord_user_id FROM verified_users WHERE uid = ?', (uid,))
result = cursor.fetchone()
return result != None
# 将验证成功的用户保存到数据库
def save_verified_user(discord_user_id: int, uid: str):
with sqlite3.connect(DATABASE_PATH) as conn:
cursor = conn.cursor()
cursor.execute('''
INSERT OR REPLACE INTO verified_users (discord_user_id, uid)
VALUES (?, ?)
''', (discord_user_id, uid))
conn.commit()
# 添加或更新用户记录
def upsert_user_record(discord_user_id: int, tid: str, zid: str, secret: str):
with sqlite3.connect(DATABASE_PATH) as conn:
cursor = conn.cursor()
cursor.execute('''
INSERT OR REPLACE INTO user_records (discord_user_id, tid, zid, secret)
VALUES (?, ?, ?, ?)
''', (discord_user_id, tid, zid, secret))
conn.commit()
# 定义一个用于 /verify 命令的交互按钮
class VerifyView(discord.ui.View):
def __init__(self, uid: str, verification_string: str, discord_user_id: int):
super().__init__(timeout=60)
self.uid = uid
self.verification_string = verification_string
self.discord_user_id = discord_user_id
@discord.ui.button(label="验证", style=discord.ButtonStyle.green)
async def verify_button(self, interaction: discord.Interaction, button: discord.ui.Button):
html_content = requests.get(f'https://www.gamemale.com/home.php?mod=space&uid={self.uid}&do=profile')
if html_content.status_code != 200:
await interaction.response.edit_message(content="请求失败,请稍后再试。", delete_after=10, view=None)
return
match = re.search(r'自定义头衔\s* </em>([^<]+)', html_content.text)
if match:
custom_title = match.group(1).strip()
if custom_title == self.verification_string:
save_verified_user(self.discord_user_id, self.uid)
await interaction.response.edit_message(content="验证成功!", delete_after=10, view=None)
else:
await interaction.response.edit_message(content="验证失败,标题不匹配。", delete_after=10, view=None)
else:
await interaction.response.edit_message(content="未找到自定义头衔,验证失败。", delete_after=10, view=None)
# 定义 /verify 命令
@bot.tree.command(name="verify", description="生成一个随机字符串,用户验证该位置")
async def verify(interaction: discord.Interaction, uid: str):
print(f'Received /verify command with uid={uid}')
# 检查用户是否已验证
is_verified, existing_uid = is_user_verified(interaction.user.id)
if is_verified:
await interaction.response.send_message("你已经验证过了!", delete_after=10, ephemeral=True)
return
if is_uid_verified(uid):
await interaction.response.send_message("该uid已经被绑定!", delete_after=10, ephemeral=True)
return
# 生成随机字符串
verification_string = pyotp.random_base32()
# 发送带有交互按钮的私密消息
await interaction.response.send_message(
content="请前往https://www.gamemale.com/home.php?mod=spacecp&ac=profile&op=info\n"
f"并在自定义头衔处填写验证字符串:{verification_string}\n"
"点击下面的按钮以确认你已填写完毕。验证完成后可任意修改。",
view=VerifyView(uid, verification_string, interaction.user.id),
delete_after=120,
ephemeral=True # 仅用户可见
)
# 定义 /addsecret 命令
@bot.tree.command(name="addsecret", description="添加或更新一个用户记录")
async def addsecret(interaction: discord.Interaction, tid: str, zid: str, secret: str):
print(f"Received /addsecret command with tid={tid}, zid={zid}, secret={secret}")
# 检查用户是否已验证
is_verified, existing_uid = is_user_verified(interaction.user.id)
if not is_verified:
await interaction.response.send_message("你还没有验证,请先进行验证。", delete_after=60, ephemeral=True)
return
# 检查是否是楼主
# TODO
# 添加或更新记录到数据库
try:
upsert_user_record(interaction.user.id, tid, zid, secret)
await interaction.response.send_message(f"记录已成功添加或更新:\nTID: {tid}\nZID: {zid}\nSecret: {secret}", delete_after=60, ephemeral=True)
except sqlite3.Error as e:
await interaction.response.send_message(f"添加或更新记录时发生错误: {e}", delete_after=60, ephemeral=True)
@bot.tree.command(name="otp", description="获取指定的 OTP 密钥")
async def otp(interaction: discord.Interaction, tid: str, zid: str):
"""处理 /otp 命令,验证用户并根据 tid 和 zid 返回 secret"""
print(f"Received /otp command with tid={tid}, zid={zid}")
# 检查用户是否已验证
is_verified, existing_uid = is_user_verified(interaction.user.id)
if not is_verified:
await interaction.response.send_message("你还没有验证,请先进行验证。", delete_after=60, ephemeral=True)
return
# 检查用户是否已购买
# TODO
# 在数据库中查找对应的 secret
with sqlite3.connect(DATABASE_PATH) as conn:
cursor = conn.cursor()
cursor.execute('''
SELECT secret FROM user_records
WHERE tid = ? AND zid = ?
''', (tid, zid))
result = cursor.fetchone()
# 检查是否找到对应的记录
if result:
db_secret = result[0]
totp = pyotp.TOTP(result[0])
TOTP = totp.now()
await interaction.response.send_message(f"验证码是{TOTP}请在30秒内输入。", delete_after=60, ephemeral=True)
else:
await interaction.response.send_message("未找到与提供的 tid 和 zid 匹配的记录。", delete_after=60, ephemeral=True)
# 启动 Bot
@bot.event
async def on_ready():
# 初始化数据库
init_db()
# 同步命令到服务器
await bot.tree.sync()
for command in bot.tree.get_commands():
print(f"Command registered: {command.name}")
print(f'Logged in as {bot.user}!')
bot.run(os.getenv('DC_TOKEN'))