feat: 添加媒体解密进度取消功能，优化检测逻辑

frontend/pages/decrypt.vue

@@ -272,7 +272,7 @@
               <div class="w-full bg-gray-200 rounded-full h-2.5 overflow-hidden">
                 <div 
                   class="h-2.5 rounded-full transition-all duration-300 ease-out"
-                  :class="decryptProgress.status === 'complete' ? 'bg-[#07C160]' : 'bg-[#91D300]'"
+                  :class="decryptProgress.status === 'complete' ? 'bg-[#07C160]' : decryptProgress.status === 'cancelled' ? 'bg-[#FAAD14]' : 'bg-[#91D300]'"
                   :style="{ width: progressPercent + '%' }"
                 ></div>
               </div>
@@ -366,6 +366,13 @@
               </svg>
               {{ mediaDecrypting ? '解密中...' : (mediaDecryptResult ? '重新解密' : '开始解密图片') }}
             </button>
+            <button
+              v-if="mediaDecrypting"
+              @click="cancelMediaDecrypt"
+              class="inline-flex items-center px-6 py-3 bg-[#FA5151] text-white rounded-lg font-medium hover:bg-[#E54D4D] transition-all duration-200"
+            >
+              停止解密
+            </button>
             <button
               @click="skipToChat"
               :disabled="mediaDecrypting"
@@ -671,6 +678,18 @@ const validateForm = () => {
 }
 
 let dbDecryptEventSource = null
+let mediaDecryptEventSource = null
+
+const closeMediaDecryptEventSource = () => {
+  try {
+    if (mediaDecryptEventSource) mediaDecryptEventSource.close()
+  } catch (e) {
+    // ignore
+  } finally {
+    mediaDecryptEventSource = null
+  }
+}
+
 onBeforeUnmount(() => {
   try {
     if (dbDecryptEventSource) dbDecryptEventSource.close()
@@ -679,6 +698,8 @@ onBeforeUnmount(() => {
   } finally {
     dbDecryptEventSource = null
   }
+
+  closeMediaDecryptEventSource()
 })
 
 const resetDbDecryptProgress = () => {
@@ -691,6 +712,17 @@ const resetDbDecryptProgress = () => {
   dbDecryptProgress.message = ''
 }
 
+const resetMediaDecryptProgress = () => {
+  decryptProgress.current = 0
+  decryptProgress.total = 0
+  decryptProgress.success_count = 0
+  decryptProgress.skip_count = 0
+  decryptProgress.fail_count = 0
+  decryptProgress.current_file = ''
+  decryptProgress.fileStatus = ''
+  decryptProgress.status = ''
+}
+
 // 处理解密
 const handleDecrypt = async () => {
   if (!validateForm()) {
@@ -877,20 +909,14 @@ const handleDecrypt = async () => {
 
 // 批量解密所有图片（使用SSE实时进度）
 const decryptAllImages = async () => {
+  closeMediaDecryptEventSource()
   mediaDecrypting.value = true
   mediaDecryptResult.value = null
   error.value = ''
   warning.value = ''
   
   // 重置进度
-  decryptProgress.current = 0
-  decryptProgress.total = 0
-  decryptProgress.success_count = 0
-  decryptProgress.skip_count = 0
-  decryptProgress.fail_count = 0
-  decryptProgress.current_file = ''
-  decryptProgress.fileStatus = ''
-  decryptProgress.status = ''
+  resetMediaDecryptProgress()
   
   try {
     // 构建SSE URL
@@ -903,8 +929,11 @@ const decryptAllImages = async () => {
     
     // 使用EventSource接收SSE
     const eventSource = new EventSource(url)
+    mediaDecryptEventSource = eventSource
     
     eventSource.onmessage = (event) => {
+      if (mediaDecryptEventSource !== eventSource) return
+
       try {
         const data = JSON.parse(event.data)
         
@@ -928,12 +957,12 @@ const decryptAllImages = async () => {
           decryptProgress.skip_count = data.skip_count
           decryptProgress.fail_count = data.fail_count
           mediaDecryptResult.value = data
-          eventSource.close()
           mediaDecrypting.value = false
+          closeMediaDecryptEventSource()
         } else if (data.type === 'error') {
           error.value = data.message
-          eventSource.close()
           mediaDecrypting.value = false
+          closeMediaDecryptEventSource()
         }
       } catch (e) {
         console.error('解析SSE消息失败:', e)
@@ -941,8 +970,10 @@ const decryptAllImages = async () => {
     }
     
     eventSource.onerror = (e) => {
+      if (mediaDecryptEventSource !== eventSource) return
+
       console.error('SSE连接错误:', e)
-      eventSource.close()
+      closeMediaDecryptEventSource()
       if (mediaDecrypting.value) {
         error.value = 'SSE连接中断，请重试'
         mediaDecrypting.value = false
@@ -951,9 +982,19 @@ const decryptAllImages = async () => {
   } catch (err) {
     error.value = err.message || '图片解密过程中发生错误'
     mediaDecrypting.value = false
+    closeMediaDecryptEventSource()
   }
 }
 
+const cancelMediaDecrypt = () => {
+  if (!mediaDecrypting.value) return
+
+  decryptProgress.status = 'cancelled'
+  mediaDecrypting.value = false
+  warning.value = '已停止图片解密，已完成的图片会保留。'
+  closeMediaDecryptEventSource()
+}
+
 // 从密钥步骤进入图片解密步骤
 const goToMediaDecryptStep = async () => {
   error.value = ''

frontend/pages/detection-result.vue

@@ -41,9 +41,19 @@
           <svg v-if="!loading" class="w-4 h-4 mr-2" fill="none" stroke="currentColor" viewBox="0 0 24 24">
             <path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M3 7v10a2 2 0 002 2h14a2 2 0 002-2V9a2 2 0 00-2-2h-6l-2-2H5a2 2 0 00-2 2z"/>
           </svg>
-          <svg v-else class="animate-spin w-4 h-4 mr-2" fill="none" viewBox="0 0 24 24">
-            <circle class="opacity-25" cx="12" cy="12" r="10" stroke="currentColor" stroke-width="4"></circle>
-            <path class="opacity-75" fill="currentColor" d="M4 12a8 8 0 018-8V0C5.373 0 0 5.373 0 12h4zm2 5.291A7.962 7.962 0 014 12H0c0 3.042 1.135 5.824 3 7.938l3-2.647z"></path>
+          <svg v-else class="w-4 h-4 mr-2 animate-spin" fill="none" viewBox="0 0 48 48" aria-hidden="true">
+            <circle class="opacity-20" cx="24" cy="24" r="18" stroke="currentColor" stroke-width="6"></circle>
+            <circle
+              cx="24"
+              cy="24"
+              r="18"
+              stroke="currentColor"
+              stroke-width="6"
+              stroke-linecap="round"
+              stroke-dasharray="28 72"
+              pathLength="100"
+              transform="rotate(-90 24 24)"
+            ></circle>
           </svg>
           {{ loading ? '检测中...' : '手动选择目录检测' }}
         </button>
@@ -53,9 +63,19 @@
       <div>
         <!-- 检测中状态 -->
         <div v-if="loading" class="absolute inset-0 bg-white/80 backdrop-blur-sm z-20 rounded-2xl flex flex-col items-center justify-center border border-[#EDEDED]">
-          <svg class="w-16 h-16 animate-spin text-[#07C160]" fill="none" stroke="currentColor" viewBox="0 0 24 24">
-            <circle class="opacity-25" cx="12" cy="12" r="10" stroke="currentColor" stroke-width="4"></circle>
-            <path class="opacity-75" fill="currentColor" d="M4 12a8 8 0 018-8V0C5.373 0 0 5.373 0 12h4zm2 5.291A7.962 7.962 0 014 12H0c0 3.042 1.135 5.824 3 7.938l3-2.647z"></path>
+          <svg class="w-16 h-16 animate-spin text-[#07C160]" fill="none" viewBox="0 0 48 48" aria-hidden="true">
+            <circle class="opacity-20" cx="24" cy="24" r="18" stroke="currentColor" stroke-width="6"></circle>
+            <circle
+              cx="24"
+              cy="24"
+              r="18"
+              stroke="currentColor"
+              stroke-width="6"
+              stroke-linecap="round"
+              stroke-dasharray="28 72"
+              pathLength="100"
+              transform="rotate(-90 24 24)"
+            ></circle>
           </svg>
           <p class="mt-4 text-lg text-[#7F7F7F]">正在检测微信数据...</p>
         </div>
@@ -395,4 +415,4 @@ onMounted(() => {
   }
   startDetection()
 })
-</script>
+</script>

src/wechat_decrypt_tool/routers/media.py

@@ -2,7 +2,7 @@ import asyncio
 import json
 from typing import Optional
 
-from fastapi import APIRouter, HTTPException
+from fastapi import APIRouter, HTTPException, Request
 from fastapi.responses import Response, StreamingResponse
 from pydantic import BaseModel, Field
 
@@ -226,6 +226,7 @@ async def get_decrypted_resource(md5: str, account: Optional[str] = None):
 
 @router.get("/api/media/decrypt_all_stream", summary="批量解密所有图片资源（SSE实时进度）")
 async def decrypt_all_media_stream(
+    request: Request,
     account: Optional[str] = None,
     xor_key: Optional[str] = None,
     aes_key: Optional[str] = None,
@@ -252,8 +253,18 @@ async def decrypt_all_media_stream(
     - 解密后非有效图片格式
     """
 
+    async def is_client_disconnected() -> bool:
+        try:
+            return await request.is_disconnected()
+        except Exception:
+            return False
+
     async def generate_progress():
         try:
+            if await is_client_disconnected():
+                logger.info("[SSE] 客户端已断开，取消图片解密任务")
+                return
+
             account_dir = _resolve_account_dir(account)
             wxid_dir = _resolve_account_wxid_dir(account_dir)
 
@@ -301,6 +312,10 @@ async def decrypt_all_media_stream(
             total_files = len(dat_files)
             logger.info(f"[SSE] 共发现 {total_files} 个.dat文件（仅图片）")
 
+            if await is_client_disconnected():
+                logger.info("[SSE] 扫描完成后客户端已断开，停止图片解密任务")
+                return
+
             if total_files == 0:
                 yield f"data: {json.dumps({'type': 'complete', 'message': '未发现需要解密的图片文件', 'total': 0, 'success_count': 0, 'skip_count': 0, 'fail_count': 0})}\n\n"
                 return
@@ -319,6 +334,10 @@ async def decrypt_all_media_stream(
             resource_dir.mkdir(parents=True, exist_ok=True)
 
             for i, (dat_path, md5) in enumerate(dat_files):
+                if await is_client_disconnected():
+                    logger.info("[SSE] 客户端已断开，停止图片解密任务")
+                    return
+
                 current = i + 1
                 file_name = dat_path.name
 

src/wechat_decrypt_tool/wechat_detection.py

@@ -16,6 +16,37 @@ from datetime import datetime
 from .database_filters import should_skip_source_database
 
 
+COMMON_WECHAT_PATTERNS = [
+    "WeChat Files",
+    "Weixin Files",
+    "wechat_files",
+    "xwechat_files",
+    "wechatMSG",
+    "WeChat",
+    "微信",
+    "Weixin",
+    "wechat",
+]
+
+SYSTEM_SCAN_SKIP_NAMES = {
+    "$recycle.bin",
+    "$winreagent",
+    "config.msi",
+    "documents and settings",
+    "intel",
+    "onedrivetemp",
+    "perflogs",
+    "program files",
+    "program files (x86)",
+    "programdata",
+    "recovery",
+    "system volume information",
+    "windows",
+    "windows.old",
+    "windows.old(1)",
+}
+
+
 def get_wx_db(msg_dir: str = None,
               db_types: Union[List[str], str] = None,
               wxids: Union[List[str], str] = None) -> List[dict]:
@@ -285,6 +316,87 @@ def get_process_list():
     return process_list
 
 
+def _is_wechat_dir_candidate_name(name: str) -> bool:
+    normalized = str(name or "").strip().lower()
+    if not normalized:
+        return False
+    return any(pattern.lower() in normalized for pattern in COMMON_WECHAT_PATTERNS)
+
+
+def _safe_iter_subdirs(directory: str) -> List[tuple[str, str]]:
+    items: List[tuple[str, str]] = []
+    try:
+        with os.scandir(directory) as entries:
+            for entry in entries:
+                try:
+                    if entry.is_dir():
+                        items.append((entry.name, entry.path))
+                except OSError:
+                    continue
+    except (PermissionError, OSError):
+        return []
+    return items
+
+
+def _append_detected_dir(detected_dirs: List[str], candidate: str) -> None:
+    if not candidate:
+        return
+    normalized = os.path.normpath(candidate)
+    if normalized not in detected_dirs:
+        detected_dirs.append(normalized)
+
+
+def _build_auto_detect_scan_paths() -> List[str]:
+    scan_paths: List[str] = []
+    seen_paths = set()
+
+    def add(path_value: str | None) -> None:
+        raw = str(path_value or "").strip()
+        if not raw:
+            return
+        normalized = os.path.normpath(raw)
+        key = normalized.lower()
+        if key in seen_paths:
+            return
+        seen_paths.add(key)
+        scan_paths.append(normalized)
+
+    home_dir = str(Path.home())
+    add(home_dir)
+    add(os.path.join(home_dir, "Documents"))
+    add(os.path.join(home_dir, "Desktop"))
+    add(os.path.join(home_dir, "Downloads"))
+
+    user_profile = str(os.environ.get("USERPROFILE") or "").strip()
+    if user_profile:
+        add(user_profile)
+        add(os.path.join(user_profile, "Documents"))
+        add(os.path.join(user_profile, "Desktop"))
+        add(os.path.join(user_profile, "Downloads"))
+
+    for drive in ("C:", "D:", "E:", "F:"):
+        drive_root = drive + os.sep
+        if not os.path.exists(drive_root):
+            continue
+
+        add(drive_root)
+
+        for child_name, child_path in _safe_iter_subdirs(drive_root):
+            if child_name.strip().lower() in SYSTEM_SCAN_SKIP_NAMES:
+                continue
+            add(child_path)
+
+        users_dir = os.path.join(drive_root, "Users")
+        add(users_dir)
+        for _user_name, user_dir in _safe_iter_subdirs(users_dir):
+            add(user_dir)
+            add(os.path.join(user_dir, "Documents"))
+            add(os.path.join(user_dir, "Desktop"))
+            add(os.path.join(user_dir, "Downloads"))
+
+    return scan_paths
+
+
 def auto_detect_wechat_data_dirs():
     """
     自动检测微信数据目录 - 多策略组合检测
@@ -292,52 +404,27 @@ def auto_detect_wechat_data_dirs():
     """
     detected_dirs = []
 
-    # 策略1：注册表检测已移除
-
-    # 策略2和策略3：注册表相关检测已移除
-
-    # 策略1：常见驱动器扫描微信相关目录
-    common_wechat_patterns = [
-        "WeChat Files", "wechat_files", "xwechat_files", "wechatMSG",
-        "WeChat", "微信", "Weixin", "wechat"
-    ]
-
-    # 扫描常见驱动器
-    drives = ['C:', 'D:', 'E:', 'F:']
-    for drive in drives:
-        if not os.path.exists(drive):
+    # 策略1：常见驱动器 / 用户目录 / 自定义目录的浅层扫描。
+    # 这里既检查扫描根目录本身，也检查其直接子目录，兼容：
+    # - C:\Users\<user>\Documents\WeChat Files
+    # - D:\wechatMSG\xwechat_files
+    # - D:\abc\wechatMSG\xwechat_files
+    for scan_path in _build_auto_detect_scan_paths():
+        if not os.path.exists(scan_path):
             continue
 
-        try:
-            # 扫描驱动器根目录和常见目录
-            scan_paths = [
-                drive + os.sep,
-                os.path.join(drive + os.sep, "Users"),
-            ]
+        scan_name = os.path.basename(os.path.normpath(scan_path))
+        if _is_wechat_dir_candidate_name(scan_name) and has_wxid_directories(scan_path):
+            _append_detected_dir(detected_dirs, scan_path)
+            print(f"[DEBUG] 目录扫描检测成功: {scan_path}")
 
-            for scan_path in scan_paths:
-                if not os.path.exists(scan_path):
-                    continue
-
-                try:
-                    for item in os.listdir(scan_path):
-                        item_path = os.path.join(scan_path, item)
-                        if not os.path.isdir(item_path):
-                            continue
-
-                        # 检查是否匹配微信目录模式
-                        for pattern in common_wechat_patterns:
-                            if pattern.lower() in item.lower():
-                                # 检查是否包含wxid目录
-                                if has_wxid_directories(item_path):
-                                    if item_path not in detected_dirs:
-                                        detected_dirs.append(item_path)
-                                        print(f"[DEBUG] 目录扫描检测成功: {item_path}")
-                                break
-                except (PermissionError, OSError):
-                    continue
-        except (PermissionError, OSError):
-            continue
+        for item_name, item_path in _safe_iter_subdirs(scan_path):
+            if not _is_wechat_dir_candidate_name(item_name):
+                continue
+            if not has_wxid_directories(item_path):
+                continue
+            _append_detected_dir(detected_dirs, item_path)
+            print(f"[DEBUG] 目录扫描检测成功: {item_path}")
 
     # 策略2：进程内存分析（简化版）
     try:
@@ -361,12 +448,11 @@ def auto_detect_wechat_data_dirs():
                             break
 
                     for parent_dir in parent_dirs:
-                        for pattern in common_wechat_patterns:
+                        for pattern in COMMON_WECHAT_PATTERNS:
                             potential_dir = os.path.join(parent_dir, pattern)
                             if os.path.exists(potential_dir) and has_wxid_directories(potential_dir):
-                                if potential_dir not in detected_dirs:
-                                    detected_dirs.append(potential_dir)
-                                    print(f"[DEBUG] 进程分析检测成功: {potential_dir}")
+                                _append_detected_dir(detected_dirs, potential_dir)
+                                print(f"[DEBUG] 进程分析检测成功: {potential_dir}")
                 except:
                     pass
     except:

tests/test_media_decrypt_stream_cancel.py

@@ -0,0 +1,74 @@
+import asyncio
+import json
+import sys
+import unittest
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from unittest import mock
+
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT / "src"))
+
+
+from wechat_decrypt_tool.routers import media as media_router  # noqa: E402  pylint: disable=wrong-import-position
+
+
+class _FakeDisconnectingRequest:
+    def __init__(self, disconnect_after: int):
+        self._disconnect_after = disconnect_after
+        self._calls = 0
+
+    async def is_disconnected(self):
+        self._calls += 1
+        return self._calls >= self._disconnect_after
+
+
+class TestMediaDecryptStreamCancel(unittest.TestCase):
+    def test_stream_stops_processing_when_client_disconnects(self):
+        with TemporaryDirectory() as td:
+            root = Path(td)
+            account_dir = root / "account"
+            wxid_dir = root / "wxid"
+            dat_path = wxid_dir / "image.dat"
+            resource_dir = account_dir / "resource"
+            wxid_dir.mkdir(parents=True, exist_ok=True)
+            dat_path.write_bytes(b"encrypted")
+
+            request = _FakeDisconnectingRequest(disconnect_after=3)
+            decrypt_mock = mock.Mock(return_value=(True, "ok"))
+
+            with mock.patch.object(media_router, "_resolve_account_dir", return_value=account_dir):
+                with mock.patch.object(media_router, "_resolve_account_wxid_dir", return_value=wxid_dir):
+                    with mock.patch.object(media_router, "_load_media_keys", return_value={"xor": 0xA5, "aes": ""}):
+                        with mock.patch.object(media_router, "_collect_all_dat_files", return_value=[(dat_path, "abc123")]):
+                            with mock.patch.object(media_router, "_get_resource_dir", return_value=resource_dir):
+                                with mock.patch.object(media_router, "_try_find_decrypted_resource", return_value=None):
+                                    with mock.patch.object(media_router, "_decrypt_and_save_resource", decrypt_mock):
+                                        response = asyncio.run(
+                                            media_router.decrypt_all_media_stream(
+                                                request=request,
+                                                account="wxid_demo",
+                                            )
+                                        )
+
+                                        async def _read_chunks():
+                                            chunks = []
+                                            async for chunk in response.body_iterator:
+                                                chunks.append(chunk.decode("utf-8") if isinstance(chunk, bytes) else str(chunk))
+                                            return chunks
+
+                                        chunks = asyncio.run(_read_chunks())
+
+            events = []
+            for chunk in chunks:
+                for line in chunk.splitlines():
+                    if line.startswith("data: "):
+                        events.append(json.loads(line[len("data: ") :]))
+
+            self.assertEqual([event.get("type") for event in events], ["scanning", "start"])
+            decrypt_mock.assert_not_called()
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_wechat_detection_auto_detect.py

@@ -0,0 +1,44 @@
+import sys
+import unittest
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from unittest.mock import patch
+
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT / "src"))
+
+
+class TestWechatDetectionAutoDetect(unittest.TestCase):
+    def test_detect_wechat_installation_finds_nested_custom_data_root(self):
+        from wechat_decrypt_tool import wechat_detection as wd
+
+        with TemporaryDirectory() as td:
+            nested_scan_root = Path(td) / "abc"
+            wechat_parent = nested_scan_root / "wechatMSG"
+            xwechat_root = wechat_parent / "xwechat_files"
+
+            login_dir = xwechat_root / "all_users" / "login" / "wxid_demo"
+            login_dir.mkdir(parents=True, exist_ok=True)
+            (login_dir / "key_info.db").write_bytes(b"demo")
+
+            account_dir = xwechat_root / "wxid_demo_nested"
+            account_dir.mkdir(parents=True, exist_ok=True)
+            (account_dir / "contact.db").write_bytes(b"demo")
+
+            with (
+                patch.object(wd, "_build_auto_detect_scan_paths", return_value=[str(nested_scan_root)]),
+                patch.object(wd, "get_process_list", return_value=[]),
+            ):
+                detected_dirs = wd.auto_detect_wechat_data_dirs()
+                result = wd.detect_wechat_installation()
+
+            self.assertEqual(detected_dirs, [str(wechat_parent)])
+            self.assertEqual(result["total_accounts"], 1)
+            self.assertEqual(result["accounts"][0]["account_name"], "wxid_demo")
+            self.assertEqual(result["accounts"][0]["data_dir"], str(account_dir))
+            self.assertEqual(result["total_databases"], 1)
+
+
+if __name__ == "__main__":
+    unittest.main()