Fix the bug that prevents the use of search, add whitelist function to prevent abuse and information leakage

Author: yym68686

README.md

@@ -18,6 +18,8 @@ Join the [Telegram Group](https://t.me/+_01cz9tAkUc1YzZl) chat to share your use
 
 ✅ 支持流式输出，实现打字机效果
 
+✅ 支持白名单，防止滥用与信息泄漏
+
 ✅ 全平台，随时随地，只要有 telegram 就可以打破知识壁垒
 
 ✅ 支持一键 Zeabur，Replit 部署，真正的零成本，傻瓜式部署，支持 kuma 防睡眠。同时支持 docker，fly.io 部署
@@ -35,6 +37,7 @@ Join the [Telegram Group](https://t.me/+_01cz9tAkUc1YzZl) chat to share your use
 | PASS_HISTORY（可选）   | 默认为真，表示机器人会记住对话历史，下次回复时会考虑上下文。如果设置为假，机器人会忘记对话历史，只考虑当前对话。 |
 | GOOGLE_API_KEY（可选） | 如果需要谷歌搜索，则需要设置。如果不设置此环境变量，机器人默认提供 duckduckgo 搜索。在 Google cloud 的 [API 与服务](https://console.cloud.google.com/apis/api/customsearch.googleapis.com) 中创建凭据，在凭据页面 API Key 就是 GOOGLE_API_KEY。Google 搜索一天可以查询 100 次，轻度使用完全足够，达到限额，机器人会自动关闭 Google 搜索。 |
 | GOOGLE_CSE_ID（可选）  | 如果需要谷歌搜索，则需要与 GOOGLE_API_KEY 一起设置。在[可编程搜索引擎](https://programmablesearchengine.google.com/) 中新建搜索引擎，其中 搜索引擎 ID 就是 GOOGLE_CSE_ID 的值。 |
+| whitelist（可选）      | 设置哪些用户可以访问机器人，将授权使用机器人的用户 ID 用`,`连接起来。默认值为`None`，即对所有人开放机器人。 |
 
 ## Zeabur 远程部署 (推荐)
 

bot.py

@@ -2,6 +2,7 @@
 import os
 import config
 import logging
+import decorators
 from md2tgmd import escape
 from runasync import run_async
 from chatgpt2api.V3 import Chatbot as GPT
@@ -26,6 +27,7 @@
 botNicKLength = len(botNick) if botNick else 0
 print("nick:", botNick)
 translator_prompt = "You are a translation engine, you can only translate text and cannot interpret it, and do not explain. Translate the text to {}, please do not explain any sentences, just translate or leave them as they are. this is the content you need to translate: "
+@decorators.Authorization
 async def command_bot(update, context, language=None, prompt=translator_prompt, title="", robot=None, has_command=True):
     if config.SEARCH_USE_GPT and not has_command:
         title = f"`🤖️ {config.DEFAULT_SEARCH_MODEL}`\n\n"
@@ -40,7 +42,7 @@ async def command_bot(update, context, language=None, prompt=translator_prompt,
                 message = prompt + message
             if config.API and message:
                 await context.bot.send_chat_action(chat_id=update.message.chat_id, action=ChatAction.TYPING)
-                await getChatGPT(update, context, title, robot, message, config.SEARCH_USE_GPT, has_command=True)
+                await getChatGPT(update, context, title, robot, message, config.SEARCH_USE_GPT, has_command)
         else:
             message = await context.bot.send_message(
                 chat_id=update.message.chat_id,
@@ -75,6 +77,7 @@ async def command_bot(update, context, language=None, prompt=translator_prompt,
         print(result)
         await context.bot.send_message(chat_id=update.message.chat_id, text=escape(result), parse_mode='MarkdownV2', disable_web_page_preview=True)
 
+@decorators.Authorization
 async def reset_chat(update, context):
     if config.API:
         config.ChatGPTbot.reset(convo_id=str(update.message.chat_id), system_prompt=config.systemprompt)
@@ -351,7 +354,7 @@ async def button_press(update, context):
             parse_mode='MarkdownV2'
         )
 
-
+@decorators.Authorization
 async def info(update, context):
     info_message = (
         f"`Hi, {update.effective_user.username}!`\n\n"
@@ -368,6 +371,7 @@ async def info(update, context):
     await context.bot.delete_message(chat_id=update.effective_chat.id, message_id=update.message.message_id)
 
 from agent import pdfQA, getmd5, persist_emdedding_pdf
+@decorators.Authorization
 async def handle_pdf(update, context):
     # 获取接收到的文件
     pdf_file = update.message.document
@@ -397,6 +401,7 @@ async def handle_pdf(update, context):
     print(result)
     await context.bot.send_message(chat_id=update.message.chat_id, text=escape(result), parse_mode='MarkdownV2', disable_web_page_preview=True)
 
+@decorators.Authorization
 async def qa(update, context):
     if (len(context.args) != 2):
         message = (
@@ -439,6 +444,7 @@ async def error(update, context):
     logger.warning('Update "%s" caused error "%s"', update, context.error)
     await context.bot.send_message(chat_id=update.message.chat_id, text="出错啦！请重试。", parse_mode='MarkdownV2')
 
+@decorators.Authorization
 async def unknown(update, context): # 当用户输入未知命令时，返回文本
     await context.bot.send_message(chat_id=update.effective_chat.id, text="Sorry, I didn't understand that command.")
 
@@ -452,9 +458,6 @@ def setup(token):
         BotCommand('zh2en', 'translate to English'),
         BotCommand('start', 'Start the bot'),
         BotCommand('reset', 'Reset the bot'),
-        # BotCommand('gpt_use_search', 'open or close gpt use search'),
-        # BotCommand('history', 'open or close chat history'),
-        # BotCommand('google', 'open or close google search'),
     ]))
 
     application.add_handler(CommandHandler("start", start))

config.py

@@ -24,4 +24,8 @@
 from chatgpt2api.V3 import Chatbot as GPT
 if API:
     ChatGPTbot = GPT(api_key=f"{API}", engine=GPT_ENGINE, system_prompt=systemprompt, temperature=temperature)
-    Claude2bot = GPT(api_key=f"{API}", engine="claude-2-web")
+    Claude2bot = GPT(api_key=f"{API}", engine="claude-2-web")
+
+whitelist = os.environ.get('whitelist', None)
+if whitelist:
+    whitelist = [int(id) for id in whitelist.split(",")]

decorators.py

@@ -0,0 +1,17 @@
+import config
+
+# 判断是否是管理员
+def Authorization(func):
+    async def wrapper(*args, **kwargs):
+        if config.whitelist == None:
+            return await func(*args, **kwargs)
+        if (args[0].effective_chat.id not in config.whitelist):
+            message = (
+                f"`Hi, {args[0].effective_user.username}!`\n\n"
+                f"id: `{args[0].effective_user.id}`\n\n"
+                f"无权访问！\n\n"
+            )
+            await args[1].bot.send_message(chat_id=args[0].effective_chat.id, text=message, parse_mode='MarkdownV2')
+            return
+        return await func(*args, **kwargs)
+    return wrapper

test/test.py

@@ -1,3 +1,3 @@
-a = "werc"
-if ("q" or "c") in a:
-    print(1)
+a = ["1", "2"]
+a = [int(i) for i in a]
+print(type(a[0]))