app.py

import streamlit as st
from pywallet import wallet  # 确保这个库能正常工作，并且你知道它的局限性
import time
import itertools
import string
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 如果密码错误, WALLET_NOT_ENCRYPTED_SENTINEL 如果未加密。
    """
    try:
        w = wallet.Wallet.from_binary(wallet_data_bytes)

        if w.is_encrypted():
            if w.decrypt(password_attempt):
                return True  # 解密成功
            else:
                return False  # 密码错误
        else:
            # 钱包未加密，返回哨兵值
            return WALLET_NOT_ENCRYPTED_SENTINEL
    except Exception as e:
        # print(f"Decryption error with password '{password_attempt}': {e}") # 调试时取消注释
        return False


def generate_passwords_bruteforce(charset, max_len):
    """
    (非常不推荐) 基本的暴力破解密码生成器
    """
    for length in range(1, max_len + 1):
        for item in itertools.product(charset, repeat=length):
            yield "".join(item)


# --- 后台工作函数 ---
def background_cracker_task(wallet_bytes_data, passwords_iterable, task_id, method_name):
    """
    在后台线程中运行破解过程。
    更新 st.session_state 来反映进度和结果。
    """
    # 初始化任务状态
    st.session_state[f'task_{task_id}_status'] = 'running'
    st.session_state[f'task_{task_id}_found_password'] = None
    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

    start_time = time.time()
    found = False
    processed_passwords = 0

    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'
            return

        processed_passwords = pwd_count + 1
        st.session_state[f'task_{task_id}_attempts'] = processed_passwords

        decryption_status = 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
        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':
        st.session_state[f'task_{task_id}_error_message'] = "未找到密码。"
        st.session_state[f'task_{task_id}_status'] = 'failed'


# --- Streamlit UI ---
st.set_page_config(page_title="Wallet.dat Cracker (Background)", layout="wide")
st.title("Bitcoin Wallet.dat 密码破解 (后台运行)")

st.warning(
    """
    **重要声明:**
    1.  此工具仅用于恢复 **您自己** 的钱包密码。
    2.  破解过程 **极其缓慢**，对复杂密码几乎无效。
    3.  **请勿将您的真实钱包文件上传到不受信任的在线服务。此工具应在本地运行。**
    4.  对于任何数据丢失或法律问题，作者概不负责。
    """
)

# 初始化 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:
    st.session_state.task_counter = 0
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()
    st.success(f"已上传 {uploaded_wallet_file.name} ({len(wallet_bytes_data_main)} bytes)")

    crack_method = st.radio(
        "2. 选择破解方法:",
        ('使用密码列表', '尝试单个密码', '基本暴力破解 (极慢，后台运行)'),
        key="crack_method_radio"
    )

    current_task_id = st.session_state.get('current_task_id')
    task_status = None
    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)
        with st.spinner(f"正在后台破解... 已尝试 {attempts} 个密码。请勿关闭此页面。"):
            st.empty()
            time.sleep(0.5)
            try:
                st.experimental_rerun()
            except AttributeError:
                st.warning("您的 Streamlit 版本较旧，spinner 状态可能不会实时更新。")


    elif task_status == 'success':
        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)

        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
            try:
                st.experimental_rerun()
            except AttributeError:
                pass


    elif task_status == 'failed':
        error_msg = st.session_state.get(f'task_{current_task_id}_error_message', "未知错误")
        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.error(f"破解失败: {error_msg}")
        st.info(f"总耗时: {time_taken:.2f} 秒。尝试密码数: {attempts}。")
        if st.button("清除结果并开始新任务", key="clear_failed"):
            st.session_state.current_task_id = None
            st.session_state.info_message = None
            try:
                st.experimental_rerun()
            except AttributeError:
                pass

    # 显示来自后台任务的一次性消息 (例如 "钱包未加密")
    # 这个消息会在 '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
        can_start = False
        start_button_label = "开始破解"
        method_name_for_task = ""

        if crack_method == '使用密码列表':
            uploaded_password_list = st.file_uploader(
                "上传密码列表文件 ( .txt, 每行一个密码)", type=["txt"], key="pwdlist_uploader"
            )
            if uploaded_password_list:
                try:
                    passwords_list = [
                        line.strip() for line in uploaded_password_list.read().decode().splitlines() if line.strip()
                    ]
                    if passwords_list:
                        st.info(f"密码列表加载了 {len(passwords_list)} 个密码。")
                        passwords_to_try_iterable = passwords_list
                        can_start = True
                        start_button_label = "开始使用列表破解"
                        method_name_for_task = "list"
                    else:
                        st.warning("上传的密码列表为空或无效。")
                except Exception as e:
                    st.error(f"读取密码列表文件错误: {e}")
            else:
                st.info("请上传密码列表文件。")


        elif crack_method == '尝试单个密码':
            single_password = st.text_input("输入要尝试的密码:", type="password", key="single_pwd_input")
            if single_password:
                passwords_to_try_iterable = [single_password]
                can_start = True
                start_button_label = "尝试此密码"
                method_name_for_task = "single"
            else:
                st.info("请输入要尝试的密码。")


        elif crack_method == '基本暴力破解 (极慢，后台运行)':
            charset_input = st.text_input("字符集:", DEFAULT_CHARSET, key="charset_input")
            max_len_input = st.number_input(
                "最大密码长度 (警告: 长度大于4会非常慢):",
                min_value=1, max_value=8, value=MAX_BRUTEFORCE_LEN, key="maxlen_input"
            )
            if charset_input:
                estimated_total = sum(len(charset_input) ** i for i in range(1, max_len_input + 1))
                st.warning(f"使用字符集 '{charset_input}' 和最大长度 {max_len_input} 进行暴力破解，"
                           f"可能需要尝试 {estimated_total} 个密码。这会非常耗时。")
                passwords_to_try_iterable = generate_passwords_bruteforce(charset_input, max_len_input)
                can_start = True
                start_button_label = "开始暴力破解 (后台)"
                method_name_for_task = "bruteforce"
            else:
                st.error("字符集不能为空。")

        if can_start and st.button(start_button_label, key=f"start_button_{method_name_for_task}"):
            st.session_state.task_counter += 1
            new_task_id = st.session_state.task_counter
            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.start()
            try:
                st.experimental_rerun()
            except AttributeError:
                pass

else:
    st.info("请先上传 wallet.dat 文件。")

st.markdown("---")
st.markdown("免责声明：此工具仅供教育和个人恢复目的。请负责任地使用。")