Upload app.py

app.py

@@ -9,15 +9,16 @@ import threading
 DEFAULT_CHARSET = string.ascii_lowercase + string.digits
 MAX_BRUTEFORCE_LEN = 4  # 保持较小以进行测试，实际破解中此长度几乎无效
 
-
 # --- 辅助函数 ---
+WALLET_NOT_ENCRYPTED_SENTINEL = "WALLET_NOT_ENCRYPTED"  # 定义一个哨兵值
+
+
 def try_decrypt_wallet(wallet_data_bytes, password_attempt):
     """
     尝试用给定的密码解密钱包。
-    返回 True 如果成功, False 如果失败。
+    返回 True 如果成功, False 如果密码错误, WALLET_NOT_ENCRYPTED_SENTINEL 如果未加密。
     """
     try:
-        # pywallet.Wallet.from_binary() 似乎是存在的，用于从字节加载
         w = wallet.Wallet.from_binary(wallet_data_bytes)
 
         if w.is_encrypted():
@@ -26,13 +27,9 @@ def try_decrypt_wallet(wallet_data_bytes, password_attempt):
             else:
                 return False  # 密码错误
         else:
-            # 如果钱包未加密，也算是一种“成功”
-            # 我们在主UI线程中通过session_state传递这个信息，而不是在后台线程直接调用st.info()
-            if 'info_message' not in st.session_state or st.session_state.info_message is None:
-                st.session_state.info_message = "钱包文件未加密。"
-            return True
+            # 钱包未加密，返回哨兵值
+            return WALLET_NOT_ENCRYPTED_SENTINEL
     except Exception as e:
-        # 实际应用中，这里可能需要更详细的错误记录
         # print(f"Decryption error with password '{password_attempt}': {e}") # 调试时取消注释
         return False
 
@@ -58,51 +55,42 @@ def background_cracker_task(wallet_bytes_data, passwords_iterable, task_id, meth
     st.session_state[f'task_{task_id}_error_message'] = None
     st.session_state[f'task_{task_id}_attempts'] = 0
     st.session_state[f'task_{task_id}_time_taken'] = 0
-    # st.session_state[f'task_{task_id}_progress'] = 0.0 # 如果可以预估总量
 
     start_time = time.time()
     found = False
     processed_passwords = 0
 
-    # 尝试获取密码总数（如果可能）
-    total_passwords_for_progress = 0
-    if isinstance(passwords_iterable, list):  # 对于列表，总数已知
-        total_passwords_for_progress = len(passwords_iterable)
-    # 对于生成器（如暴力破解），总数通常未知或难以计算，除非预先计算
-
     for pwd_count, pwd in enumerate(passwords_iterable):
-        # 检查任务是否被中止 (例如，用户启动了新任务或关闭了应用)
         if st.session_state.get('current_task_id') != task_id:
             st.session_state[f'task_{task_id}_error_message'] = "任务被中止。"
-            st.session_state[f'task_{task_id}_status'] = 'failed'  # 或 'aborted'
+            st.session_state[f'task_{task_id}_status'] = 'failed'
             return
 
         processed_passwords = pwd_count + 1
         st.session_state[f'task_{task_id}_attempts'] = processed_passwords
 
-        # 更新进度（可选，如果spinner的文本不直接显示这个）
-        # if total_passwords_for_progress > 0:
-        #     st.session_state[f'task_{task_id}_progress'] = processed_passwords / total_passwords_for_progress
+        decryption_status = try_decrypt_wallet(wallet_bytes_data, pwd)
 
-        if try_decrypt_wallet(wallet_bytes_data, pwd):
+        if decryption_status is True:  # 密码正确
             st.session_state[f'task_{task_id}_found_password'] = pwd
             st.session_state[f'task_{task_id}_status'] = 'success'
             found = True
             break
-
-        # 对于非常快的循环，可以考虑不那么频繁地更新session_state以减少开销
-        # 例如: if processed_passwords % 100 == 0: st.session_state...
-        # 但目前每次更新是为了让spinner中的计数尽可能准确
+        elif decryption_status == WALLET_NOT_ENCRYPTED_SENTINEL:  # 钱包未加密
+            st.session_state.info_message = "钱包文件未加密。"  # 在这里设置
+            st.session_state[f'task_{task_id}_found_password'] = WALLET_NOT_ENCRYPTED_SENTINEL  # 使用哨兵值标记
+            st.session_state[f'task_{task_id}_status'] = 'success'  # 同样视为成功
+            found = True
+            break
+        # else: decryption_status is False, 密码错误，继续循环
 
     end_time = time.time()
     st.session_state[f'task_{task_id}_time_taken'] = end_time - start_time
 
-    if not found and st.session_state.get(f'task_{task_id}_status') == 'running':  # 确保不是因为中止而退出
+    if not found and st.session_state.get(f'task_{task_id}_status') == 'running':
         st.session_state[f'task_{task_id}_error_message'] = "未找到密码。"
         st.session_state[f'task_{task_id}_status'] = 'failed'
 
-    # 线程结束后，Streamlit 会因为 session_state 的变化在下一次 rerun 时更新 UI
-
 
 # --- Streamlit UI ---
 st.set_page_config(page_title="Wallet.dat Cracker (Background)", layout="wide")
@@ -121,23 +109,22 @@ st.warning(
 # 初始化 session_state 变量 (如果它们不存在)
 if 'current_task_id' not in st.session_state:
     st.session_state.current_task_id = None
-if 'task_counter' not in st.session_state:  # 用于生成唯一的 task_id
+if 'task_counter' not in st.session_state:
     st.session_state.task_counter = 0
-if 'info_message' not in st.session_state:  # 用于一次性消息，如“钱包未加密”
+if 'info_message' not in st.session_state:
     st.session_state.info_message = None
 
 # --- 文件上传 ---
 uploaded_wallet_file = st.file_uploader("1. 上传你的 wallet.dat 文件", type=["dat"])
 
 if uploaded_wallet_file:
-    wallet_bytes_data_main = uploaded_wallet_file.getvalue()  # 在主线程中读取文件内容
+    wallet_bytes_data_main = uploaded_wallet_file.getvalue()
     st.success(f"已上传 {uploaded_wallet_file.name} ({len(wallet_bytes_data_main)} bytes)")
 
-    # --- 破解方法选择 ---
     crack_method = st.radio(
         "2. 选择破解方法:",
         ('使用密码列表', '尝试单个密码', '基本暴力破解 (极慢，后台运行)'),
-        key="crack_method_radio"  # 添加key以帮助Streamlit跟踪小部件状态
+        key="crack_method_radio"
     )
 
     current_task_id = st.session_state.get('current_task_id')
@@ -145,30 +132,34 @@ if uploaded_wallet_file:
     if current_task_id:
         task_status = st.session_state.get(f'task_{current_task_id}_status')
 
-    # --- 显示任务状态和结果 ---
     if task_status == 'running':
         attempts = st.session_state.get(f'task_{current_task_id}_attempts', 0)
-        # progress_val = st.session_state.get(f'task_{current_task_id}_progress', 0.0)
         with st.spinner(f"正在后台破解... 已尝试 {attempts} 个密码。请勿关闭此页面。"):
-            # st.progress(progress_val) # 如果有进度条
-            st.empty()  # 保持 spinner 可见
-            time.sleep(0.5)  # 调整此处的休眠时间，控制UI刷新频率
+            st.empty()
+            time.sleep(0.5)
             try:
                 st.experimental_rerun()
-            except AttributeError:  # 兼容非常旧的 Streamlit 版本
+            except AttributeError:
                 st.warning("您的 Streamlit 版本较旧，spinner 状态可能不会实时更新。")
 
 
     elif task_status == 'success':
-        found_pwd = st.session_state.get(f'task_{current_task_id}_found_password')
+        found_pwd_or_status = st.session_state.get(f'task_{current_task_id}_found_password')
         time_taken = st.session_state.get(f'task_{current_task_id}_time_taken', 0)
         attempts = st.session_state.get(f'task_{current_task_id}_attempts', 0)
-        st.success(f"🎉 密码找到: **{found_pwd}** 🎉")
-        st.balloons()
+
+        if found_pwd_or_status == WALLET_NOT_ENCRYPTED_SENTINEL:
+            st.success(f"✅ 钱包文件未加密。")
+            # info_message 应该已经被后台任务设置，这里不需要再次设置
+        else:
+            st.success(f"🎉 密码找到: **{found_pwd_or_status}** 🎉")
+            st.balloons()
+
         st.info(f"总耗时: {time_taken:.2f} 秒。尝试密码数: {attempts}。")
+
         if st.button("清除结果并开始新任务", key="clear_success"):
             st.session_state.current_task_id = None
-            st.session_state.info_message = None  # 清除可能存在的 "钱包未加密" 消息
+            st.session_state.info_message = None
             try:
                 st.experimental_rerun()
             except AttributeError:
@@ -189,14 +180,16 @@ if uploaded_wallet_file:
             except AttributeError:
                 pass
 
-    # 显示来自 try_decrypt_wallet 的一次性消息 (例如 "钱包未加密")
-    if st.session_state.get('info_message') and task_status != 'running':  # 仅当任务不在运行时显示
-        st.info(st.session_state.info_message)
-        # 不立即清除，允许在结果旁边显示。清除按钮会处理它。
+    # 显示来自后台任务的一次性消息 (例如 "钱包未加密")
+    # 这个消息会在 'success' 状态时由后台任务设置
+    # 确保它只在相关任务完成后显示一次，并且在清除任务时也被清除
+    if st.session_state.get('info_message') and task_status == 'success' and st.session_state.get(
+            f'task_{current_task_id}_found_password') == WALLET_NOT_ENCRYPTED_SENTINEL:
+        # st.info(st.session_state.info_message) # 这行其实可以省略，因为 "钱包文件未加密" 已经在 success 块中显示
+        pass  # The message is handled within the success block now
 
-    # --- 启动按钮 (仅当没有任务运行时显示或任务已完成/失败) ---
     if not task_status or task_status in ['success', 'failed']:
-        passwords_to_try_iterable = None  # 使用可迭代对象
+        passwords_to_try_iterable = None
         can_start = False
         start_button_label = "开始破解"
         method_name_for_task = ""
@@ -207,7 +200,6 @@ if uploaded_wallet_file:
             )
             if uploaded_password_list:
                 try:
-                    # 直接将文件对象传递或读取为列表
                     passwords_list = [
                         line.strip() for line in uploaded_password_list.read().decode().splitlines() if line.strip()
                     ]
@@ -228,7 +220,7 @@ if uploaded_wallet_file:
         elif crack_method == '尝试单个密码':
             single_password = st.text_input("输入要尝试的密码:", type="password", key="single_pwd_input")
             if single_password:
-                passwords_to_try_iterable = [single_password]  # 包装在列表中
+                passwords_to_try_iterable = [single_password]
                 can_start = True
                 start_button_label = "尝试此密码"
                 method_name_for_task = "single"
@@ -259,17 +251,16 @@ if uploaded_wallet_file:
             st.session_state.current_task_id = new_task_id
             st.session_state.info_message = None  # 清除旧的 "钱包未加密" 消息，准备新任务
 
-            # 启动后台线程
             thread = threading.Thread(
                 target=background_cracker_task,
                 args=(wallet_bytes_data_main, passwords_to_try_iterable, new_task_id, method_name_for_task)
             )
-            thread.daemon = True  # 确保主程序退出时线程也退出
+            thread.daemon = True
             thread.start()
             try:
-                st.experimental_rerun()  # 立即重新运行脚本以显示 spinner
+                st.experimental_rerun()
             except AttributeError:
-                pass  # 旧版 Streamlit
+                pass
 
 else:
     st.info("请先上传 wallet.dat 文件。")