v0.0.10

GenshinMinecraft · Jun 25, 2024 · ab3b1dc · ab3b1dc
1 parent 6599c83
commit ab3b1dc
Show file tree

Hide file tree

Showing 6 changed files with 164 additions and 44 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "Rust-Cloudflare-Workers-AI-Telegram-Bot"
-version = "0.0.9"
+version = "0.0.10"
 edition = "2021"
 
 [dependencies]
@@ -17,7 +17,6 @@ lazy_static = "1.4.0"
 # Must always use panic = "abort" to avoid needing to define the unstable eh_personality lang item.
 panic = "abort"
 
-
 [profile.release]
 opt-level = "z"     # Optimize for size.
 lto = true          # Enable Link Time Optimization

diff --git a/README.md b/README.md
@@ -1,63 +1,66 @@
 # Rust-Cloudflare-Workers-AI-Telegram-Bot
 
-其 Python 实现请看 [Cloudflare-Workers-Ai-Telegram-Bot](https://github.com/GenshinMinecraft/Cloudflare-Workers-Ai-Telegram-Bot)
+[简体中文](https://github.com/GenshinMinecraft/Rust-Cloudflare-Workers-AI-Telegram-Bot/blob/main/README_cn.md)
 
-## 简介
+For its Python implementation, please refer to [Cloudflare-Workers-Ai-Telegram-Bot](https://github.com/GenshinMinecraft/Cloudflare-Workers-Ai-Telegram-Bot).
 
-一个使用 Rust 写的 Cloudflare Workers AI Bot，支持**上下文**
+## Introduction
 
-通过调用 Workers AI REST API 来与用户对话
-
-通过 Workers KV 储存用户的上下文信息
+This is a Rust-based Cloudflare Workers AI Bot with support for **context preservation**. It interacts with users through invoking the Workers AI REST API and stores user context information using Workers KV.
 
 ## To Do List
 
-- [x] 上下文支持
-- [x] 代理连接 (通过 reqwest 本身支持的环境变量实现)
-- [ ] 用户/群组 ID 鉴权
+- [x] Context preservation implemented
+- [x] Proxy support via environment variables (leveraging reqwest's native support)
+- [x] User/Group ID authentication
+
+For additional features, please raise issues or submit pull requests.
 
-## 运行
+## Running the Bot
 
-在 Release 中下载并配置好可执行权限
+Download the executable from Releases and set execution permissions accordingly.
 
-请在环境中使用**变量**填写好信息后 使用
+Configure the required **environment variables** before running:
 
-- `API_KEY` (必须): 
-    在[这里](https://dash.cloudflare.com/profile/api-tokens)获取，至少需要 Workers AI 的读写和 Workers KV 的读写权限
-- `USER_ID` (必须): 
-    Cloudflare 的 Account ID，最简单的获取方式就是打开 Cloudflare Dash，URL 中的那串就是，比如 `41810b51b9f7521da5fea96d12xxxxxx`
-- `PROMPT` (非必须): 
-    AI 提示词
-- `MODEL` (非必须): 
-    对话使用的大模型，默认是阿里云的通义千问，可以在[这里](https://developers.cloudflare.com/workers-ai/models/)查看支持的模型
-- `KV_NAMESPACE_ID` (必须): 
-    Workers KV 的 Namespace ID
-- `TELEGRAM_BOTTOKEN` (必须): 
-    TG BOT TOKEN
-- `http_proxy` 与 `https_proxy` (非必须): 
-    这是 Linux Bash 环境下通用的变量，可以填写 http(s) 或 socks 代理来请求 Telegram 与 Cloudflare API
+- `API_KEY` (mandatory): 
+    Obtain from [here](https://dash.cloudflare.com/profile/api-tokens), requiring read/write access to Workers AI and Workers KV.
+- `USER_ID` (mandatory): 
+    Your Cloudflare Account ID, easily found in the URL when accessing Cloudflare Dashboard, e.g., `41810b51b9f7521da5fea96d12xxxxxx`.
+- `PROMPT` (optional): 
+    Custom AI prompt.
+- `MODEL` (optional): 
+    The large language model for conversations, defaulting to Alibaba Cloud's Qwen. View supported models [here](https://developers.cloudflare.com/workers-ai/models/).
+- `KV_NAMESPACE_ID` (mandatory): 
+    The ID of your Workers KV Namespace.
+- `TELEGRAM_BOTTOKEN` (mandatory): 
+    Your Telegram Bot token.
+- `TELEGRAM_ID` (optional): 
+    Whitelist Telegram IDs; only these groups/users can interact with the bot. Separate multiple IDs with `,`.
+- `http_proxy` & `https_proxy` (optional): 
+    Standard Bash variables for proxy configuration, allowing HTTP(S) or SOCKS proxies for accessing Telegram and Cloudflare APIs.
 
-比如 Linux Bash 环境下:
+In a Linux Bash environment, run as follows:
 
 ```bash
-API_KEY="xxx" \
-USER_ID="xxx" \
-PROMPT="xxx" \
-MODEL="xxx" \
-KV_NAMESPACE_ID="xxx" \
-TELEGRAM_BOTTOKEN="xxx" \
+API_KEY="YOUR_SECRET_KEY" \
+USER_ID="YOUR_ACCOUNT_ID" \
+PROMPT="You Are A Pig" \
+MODEL="@cf/qwen/qwen1.5-14b-chat-awq" \
+KV_NAMESPACE_ID="NAMESPACE_ID" \
+TELEGRAM_BOTTOKEN="YOUR_BOT_TOKEN" \
+TELEGRAM_ID="123,-10086,114514" \
 ./amd64-linux
 ```
 
-运行即可，Bot 信息请看 Log 输出
+The bot will start, and log outputs will provide operational details.
 
-## 鸣谢
+## Acknowledgments
 
-- [Cloudflare](https://cloudflare.com): 赞美大爹
-- [Teloxide](https://github.com/teloxide/teloxide): Bot 主框架
+- [Cloudflare](https://cloudflare.com): Kudos to the mighty provider.
+- [Teloxide](https://github.com/teloxide/teloxide): The backbone framework for the bot.
 
 ## LICENSE
 
-本项目采用 **你他妈的想干嘛就干嘛公共许可证** *(WTFPL)*
+This project is licensed under the **Do What the Fuck You Want to Public License** (WTFPL).
 
-所以，你可以对这个项目干任何你想干的事情 *(你他妈的想干嘛就干嘛)*
+So, feel free to do whatever the hell you want with this project—literally, *do what the fuck you want*.
diff --git a/README_cn.md b/README_cn.md
@@ -0,0 +1,68 @@
+# Rust-Cloudflare-Workers-AI-Telegram-Bot
+
+其 Python 实现请看 [Cloudflare-Workers-Ai-Telegram-Bot](https://github.com/GenshinMinecraft/Cloudflare-Workers-Ai-Telegram-Bot)
+
+## 简介
+
+一个使用 Rust 写的 Cloudflare Workers AI Bot，支持**上下文**
+
+通过调用 Workers AI REST API 来与用户对话
+
+通过 Workers KV 储存用户的上下文信息
+
+## To Do List
+
+- [x] 上下文支持
+- [x] 代理连接 (通过 reqwest 本身支持的环境变量实现)
+- [x] 用户/群组 ID 鉴权
+
+需要更多功能请提出 Issue 或提交 PR
+
+## 运行
+
+在 Release 中下载并配置好可执行权限
+
+请在环境中使用**变量**填写好信息后 使用
+
+- `API_KEY` (必须): 
+    在[这里](https://dash.cloudflare.com/profile/api-tokens)获取，至少需要 Workers AI 的读写和 Workers KV 的读写权限
+- `USER_ID` (必须): 
+    Cloudflare 的 Account ID，最简单的获取方式就是打开 Cloudflare Dash，URL 中的那串就是，比如 `41810b51b9f7521da5fea96d12xxxxxx`
+- `PROMPT` (非必须): 
+    AI 提示词
+- `MODEL` (非必须): 
+    对话使用的大模型，默认是阿里云的通义千问，可以在[这里](https://developers.cloudflare.com/workers-ai/models/)查看支持的模型
+- `KV_NAMESPACE_ID` (必须): 
+    Workers KV 的 Namespace ID
+- `TELEGRAM_BOTTOKEN` (必须): 
+    TG BOT TOKEN
+- `TELEGRAM_ID` (非必须):
+    Telegram 白名单 ID，只有白名单内的群组/用户才可使用，多个 ID 使用 `,`(英文逗号) 隔开
+- `http_proxy` 与 `https_proxy` (非必须): 
+    这是 Linux Bash 环境下通用的变量，可以填写 http(s) 或 socks 代理来请求 Telegram 与 Cloudflare API
+
+比如 Linux Bash 环境下:
+
+```bash
+API_KEY="SECRET" \
+USER_ID="SECRET" \
+PROMPT="You Are A Pig" \
+MODEL="@cf/qwen/qwen1.5-14b-chat-awq" \
+KV_NAMESPACE_ID="SECRET" \
+TELEGRAM_BOTTOKEN="SECRET" \
+TELEGRAM_ID="123,-10086,114514" \
+./amd64-linux
+```
+
+运行即可，Bot 信息请看 Log 输出
+
+## 鸣谢
+
+- [Cloudflare](https://cloudflare.com): 赞美大爹
+- [Teloxide](https://github.com/teloxide/teloxide): Bot 主框架
+
+## LICENSE
+
+本项目采用 **你他妈的想干嘛就干嘛公共许可证** *(WTFPL)*
+
+所以，你可以对这个项目干任何你想干的事情 *(你他妈的想干嘛就干嘛)*
diff --git a/src/config.rs b/src/config.rs
@@ -12,6 +12,7 @@ pub struct Config {
     pub model: String,
     pub kv_namespace_id: String,
     pub telegram_bottoken: String,
+    pub telegram_id: Vec<i128>,
 }
 
 pub fn setup_config() -> Config {
@@ -22,6 +23,7 @@ pub fn setup_config() -> Config {
         model: String::new(),
         kv_namespace_id: String::new(),
         telegram_bottoken: String::new(),
+        telegram_id: Vec::new(),
     };
 
     match env::var("API_KEY") {
@@ -90,5 +92,25 @@ pub fn setup_config() -> Config {
         },
     };
 
+    match env::var("TELEGRAM_ID") {
+        Ok(tmp) => {
+            let array: Vec<i128> = if tmp.contains(',') {
+                // 按逗号分割
+                tmp.split(',')
+                    .map(|s| s.trim().parse().expect("无法解析为数字"))
+                    .collect()
+            } else {
+                // 将整个字符串解析为一个数字
+                vec![tmp.trim().parse().expect("无法解析为数字")]
+            };
+
+            config.telegram_id = array;
+        },
+        Err(_) => {
+            warn!("无法获取 TELEGRAM_ID 的值，将开放给所有用户使用");
+            config.telegram_id = vec![1145141919810]; // 1145141919810 代表所有用户 (不要问为什么这样设置，问就是喜欢)
+        },
+    }
+
     return config;
 }
diff --git a/src/main.rs b/src/main.rs
@@ -76,6 +76,32 @@ async fn reply_start(msg: Message, bot: Bot) {
 
 // 处理 ai 回复
 async fn reply_ai(msg: Message, bot: Bot, optiontext: Option<&str>) {
+
+    // 检测是否为可使用的用户
+    if CFG.telegram_id == vec![1145141919810] {
+        // 无白名单直接使用
+    } else {
+        if msg.chat.is_private() { // 私聊
+            let is_user_available: bool = CFG.telegram_id.iter().any(|&x| x == msg.chat.id.to_string().parse().unwrap());
+            if is_user_available {
+                // 在白名单内
+            } else { // 不在白名单内
+                let _ = bot.send_message(msg.chat.id, "❌ 您不在白名单内，请联系管理员授权").await;
+                info!("用户未授权");
+                return;
+            }
+        } else { // 群组
+            let is_group_available: bool = CFG.telegram_id.iter().any(|&x| x == msg.chat.id.to_string().parse().unwrap());
+            if is_group_available {
+                // 在白名单内
+            } else { // 不在白名单内
+                let _ = bot.send_message(msg.chat.id, "❌ 该群组不在白名单内，请联系管理员授权").await;
+                info!("用户未授权");
+                return;
+            }
+        }
+    }
+
     let text: &str;
 
     // 检测是否有参数
@@ -84,7 +110,7 @@ async fn reply_ai(msg: Message, bot: Bot, optiontext: Option<&str>) {
         None => {
             warn!("{}", format!("用户 {:?} 使用方法不正确", get_user_id(msg.clone()).await.unwrap_or("无法获取 ID".to_string())));
 
-            let _ = bot.send_message(msg.chat.id, "使用方法不正确！请使用 /start 来查看使用方法")
+            let _ = bot.send_message(msg.chat.id, "❌ 使用方法不正确！请使用 /start 来查看使用方法")
             .parse_mode(ParseMode::MarkdownV2)
             .await;
             return;
@@ -174,13 +200,14 @@ async fn reply_clear(msg: Message, bot: Bot) {
         Err(msg) => error!("{}", msg.to_string()),
     }
 
-    let _ = bot.send_message(msg.chat.id, "成功清除上下文")
+    let _ = bot.send_message(msg.chat.id, "✔️ 成功清除上下文")
     .parse_mode(ParseMode::MarkdownV2)
     .await;
 }
 
 // 私聊处理
 async fn match_msg_private(msg: Message, bot: Bot) {
+
     let text: &str = msg.text().unwrap();
 
     if text.starts_with('/') { // 是否为 "/" 开头的命令
@@ -204,6 +231,7 @@ async fn match_msg_private(msg: Message, bot: Bot) {
 
 // 非私聊处理
 async fn match_msg_public(msg: Message, bot: Bot) {
+
     let text: &str = msg.text().unwrap();
 
     if text.starts_with('/') { // 是否为 "/" 开头的命令