From 214dadad45d59c6ef821a93f3ca023d401a9d0db Mon Sep 17 00:00:00 2001
From: Xiaolan Bot <xiaolan@local>
Date: Wed, 25 Feb 2026 15:21:47 +0800
Subject: [PATCH] feat: add /update command for owner-only self-update and
 restart

---
 .env.example |  3 ++
 README.md    |  4 +++
 SubMind.py   | 82 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 89 insertions(+)

diff --git a/.env.example b/.env.example
index 7f99c4d..b073c91 100644
--- a/.env.example
+++ b/.env.example
@@ -1,2 +1,5 @@
 TELEGRAM_TOKEN=""
 EXCHANGE_API_KEY=""
+UPDATE_OWNER_ID=""
+AUTO_UPDATE_REMOTE="gitllc"
+AUTO_UPDATE_BRANCH="main"
diff --git a/README.md b/README.md
index ea1b204..197dde6 100644
--- a/README.md
+++ b/README.md
@@ -61,6 +61,9 @@ EXCHANGE_API_KEY="<YOUR_EXCHANGE_API_KEY>"
 说明：
 - `TELEGRAM_TOKEN` 必填。
 - `EXCHANGE_API_KEY` 可选（不填时不做在线汇率转换）。
+- `UPDATE_OWNER_ID` 可选（建议配置为你的 Telegram 用户 ID，仅该用户可执行 `/update`）。
+- `AUTO_UPDATE_REMOTE` 可选（默认 `gitllc`）。
+- `AUTO_UPDATE_BRANCH` 可选（默认 `main`）。
 
 ### 4) 运行
 
@@ -80,6 +83,7 @@ python SubMind.py
 - `/import` 导入 CSV
 - `/export` 导出 CSV
 - `/set_currency <CODE>` 设置主货币（例如 `USD`、`CNY`）
+- `/update` 拉取最新代码、安装依赖并自动重启（仅 `UPDATE_OWNER_ID` 指定用户可用）
 - `/help` 帮助
 - `/cancel` 取消当前流程
 
diff --git a/SubMind.py b/SubMind.py
index a2433b3..e4092a2 100644
--- a/SubMind.py
+++ b/SubMind.py
@@ -1,6 +1,8 @@
 import sqlite3
 import asyncio
 import os
+import sys
+import subprocess
 import html
 import requests
 import datetime
@@ -44,6 +46,11 @@ EXCHANGE_API_KEY = os.getenv('EXCHANGE_API_KEY')
 PROJECT_NAME = "SubMind"
 DB_FILE = 'submind.db'
 
+# 自动更新配置
+UPDATE_OWNER_ID = os.getenv('UPDATE_OWNER_ID')  # 仅允许此用户执行 /update
+AUTO_UPDATE_REMOTE = os.getenv('AUTO_UPDATE_REMOTE', 'gitllc')
+AUTO_UPDATE_BRANCH = os.getenv('AUTO_UPDATE_BRANCH', 'main')
+
 # --- 对话处理器状态 ---
 (ADD_NAME, ADD_COST, ADD_CURRENCY, ADD_CATEGORY, ADD_NEXT_DUE,
  ADD_FREQ_UNIT, ADD_FREQ_VALUE, ADD_RENEWAL_TYPE, ADD_NOTES) = range(9)
@@ -1608,6 +1615,79 @@ async def cancel(update: Update, context: CallbackContext):
     return ConversationHandler.END
 
 
+def _can_run_update(user_id: int) -> bool:
+    """仅允许指定 owner 执行自动更新。未配置 owner 时默认拒绝。"""
+    if not UPDATE_OWNER_ID:
+        return False
+    try:
+        return int(UPDATE_OWNER_ID) == int(user_id)
+    except (ValueError, TypeError):
+        return False
+
+
+async def update_bot(update: Update, context: CallbackContext):
+    user_id = update.effective_user.id
+    if not _can_run_update(user_id):
+        await update.message.reply_text("无权限执行 /update。")
+        return
+
+    await update.message.reply_text("开始检查更新，请稍候…")
+
+    repo_dir = os.path.dirname(os.path.abspath(__file__))
+
+    try:
+        fetch_cmd = ["git", "fetch", AUTO_UPDATE_REMOTE, AUTO_UPDATE_BRANCH]
+        fetch_proc = subprocess.run(fetch_cmd, cwd=repo_dir, capture_output=True, text=True)
+        if fetch_proc.returncode != 0:
+            err = (fetch_proc.stderr or fetch_proc.stdout or "未知错误").strip()
+            await update.message.reply_text(f"更新失败（fetch）：\n<code>{escape_html(err)}</code>", parse_mode='HTML')
+            return
+
+        local_rev = subprocess.run(
+            ["git", "rev-parse", "HEAD"], cwd=repo_dir, capture_output=True, text=True
+        )
+        remote_rev = subprocess.run(
+            ["git", "rev-parse", f"{AUTO_UPDATE_REMOTE}/{AUTO_UPDATE_BRANCH}"],
+            cwd=repo_dir, capture_output=True, text=True
+        )
+
+        if local_rev.returncode != 0 or remote_rev.returncode != 0:
+            await update.message.reply_text("更新失败：无法读取当前版本。")
+            return
+
+        local_hash = local_rev.stdout.strip()
+        remote_hash = remote_rev.stdout.strip()
+
+        if local_hash == remote_hash:
+            await update.message.reply_text("当前已是最新版本，无需更新。")
+            return
+
+        reset_proc = subprocess.run(
+            ["git", "reset", "--hard", f"{AUTO_UPDATE_REMOTE}/{AUTO_UPDATE_BRANCH}"],
+            cwd=repo_dir, capture_output=True, text=True
+        )
+        if reset_proc.returncode != 0:
+            err = (reset_proc.stderr or reset_proc.stdout or "未知错误").strip()
+            await update.message.reply_text(f"更新失败（reset）：\n<code>{escape_html(err)}</code>", parse_mode='HTML')
+            return
+
+        pip_proc = subprocess.run(
+            [sys.executable, "-m", "pip", "install", "-r", "requirements.txt"],
+            cwd=repo_dir, capture_output=True, text=True
+        )
+        if pip_proc.returncode != 0:
+            err = (pip_proc.stderr or pip_proc.stdout or "未知错误").strip()
+            await update.message.reply_text(f"依赖安装失败：\n<code>{escape_html(err[-1800:])}</code>", parse_mode='HTML')
+            return
+
+        await update.message.reply_text("更新完成，正在重启机器人…")
+        os.execv(sys.executable, [sys.executable] + sys.argv)
+
+    except Exception as e:
+        logger.error(f"/update failed: {e}")
+        await update.message.reply_text(f"更新异常：<code>{escape_html(str(e))}</code>", parse_mode='HTML')
+
+
 # --- Main ---
 def main():
     if not TELEGRAM_TOKEN:
@@ -1636,6 +1716,7 @@ def main():
             BotCommand("import", "📥 导入订阅"),
             BotCommand("export", "📤 导出订阅"),
             BotCommand("set_currency", "💲 设置主货币"),
+            BotCommand("update", "🛠️ 拉取最新代码并重启"),
             BotCommand("help", "ℹ️ 获取帮助"),
             BotCommand("cancel", "❌ 取消当前操作")
         ]
@@ -1730,6 +1811,7 @@ def main():
     application.add_handler(CommandHandler('set_currency', set_currency))
     application.add_handler(CommandHandler('stats', stats))
     application.add_handler(CommandHandler('export', export_command))
+    application.add_handler(CommandHandler('update', update_bot))
     application.add_handler(CommandHandler('cancel', cancel))
 
     application.add_handler(add_conv)