import sys
import urllib.parse
import base64
import re
import json
import os
from mitmproxy import http
from mitmproxy.tools.main import mitmdump

# --- 配置 ---
TARGET_URL = "http://baodao.*.edu.cn/XSSC/Question.aspx?tid="
JSON_FILE = "question_bank.json"
# DEBUG_LOG_FILE = "viewstate_dump.txt" # 调试日志文件 (已禁用)
# --- 配置结束 ---

# def log_debug(message): # 调试日志函数 (已禁用)
#     """打印调试信息"""
#     # print(f"[DEBUG] {message}")
#     pass

# def dump_to_file(content): # 转储函数 (已禁用)
#     """将原始 ViewState 写入文件以便分析"""
#     # try:
#     #     with open(DEBUG_LOG_FILE, 'a', encoding='utf-8') as f:
#     #         f.write(content + "\n" + "="*20 + "\n")
#     # except Exception as e:
#     #     print(f"[ERROR] 无法写入调试文件: {e}")
#     pass

def load_question_bank():
    if os.path.exists(JSON_FILE):
        try:
            with open(JSON_FILE, 'r', encoding='utf-8') as f:
                # 使用 set 来存储已知题目的 question 字符串，以便快速查找
                bank_data = json.load(f)
                known_questions = set(bank_data.keys())
                return bank_data, known_questions
        except json.JSONDecodeError:
            print(f"[ERROR] {JSON_FILE} 文件损坏, 将创建新文件。")
            return {}, set()
        except Exception as e:
            print(f"[ERROR] 加载题库时出错: {e}")
            return {}, set()
    return {}, set()

def save_question_bank(bank):
    try:
        with open(JSON_FILE, 'w', encoding='utf-8') as f:
            json.dump(bank, f, ensure_ascii=False, indent=2)
    except IOError as e:
        print(f"[ERROR] 无法写入题库文件: {e}")

class ViewStateSniffer:
    def __init__(self):
        self.bank, self.known_questions = load_question_bank()
        print(f"[INFO] 启动 ViewState 嗅探器... 已加载 {len(self.bank)} 道题目。")
        print(f"[INFO] 拦截目标: {TARGET_URL}")
        print(f"[INFO] 题库文件: {JSON_FILE}")
        # print(f"[INFO] 调试日志: {DEBUG_LOG_FILE} (原始 ViewState 将被转储于此)") # 已禁用
        # 清空上次的调试日志 (已禁用)
        # if os.path.exists(DEBUG_LOG_FILE):
        #     os.remove(DEBUG_LOG_FILE)

    def parse_viewstate_bytes(self, data: bytes):
        # log_debug("开始解析 ViewState 二进制数据...")
        try:
            # 1. 解码为 UTF-8 文本
            text = data.decode('utf-8', errors='ignore')
            # log_debug(f"UTF-8 解码后文本 (前 500 字符): {text[:500]}")

            # 2. 提取题目
            q_match = re.search(r'(\d+\.\s.*?（\s*）)', text)
            if not q_match:
                q_match = re.search(r'(\d+\..*?\([^\)]*\))', text) # 兼容英文括号

            if not q_match:
                # log_debug("[PARSE_FAIL] 步骤 2 失败: 未找到题目")
                return None, [], None

            question = q_match.group(1).strip()
            # log_debug(f"[PARSE_OK] 步骤 2 成功: 找到题目: {question}")

            # 3. 提取选项
            options_text = text[q_match.end():]
            # log_debug(f"开始在以下文本中查找选项 (前 300 字符): {options_text[:300]}")
            options_matches = re.findall(r'([A-D]\.[^\x00#,<]+)', options_text)
            options = [opt.strip().rstrip(',').rstrip('#').strip() for opt in options_matches][:4]

            if len(options) != 4:
                # log_debug(f"[PARSE_FAIL] 步骤 3 失败: 找到 {len(options)} 个选项，不等于 4。")
                # log_debug(f"找到的选项: {options}")
                # log_debug("尝试后备选项提取逻辑...")
                options = []
                temp_options_text = text[q_match.end():]
                for opt_char in ['A', 'B', 'C', 'D']:
                    opt_match = re.search(rf'({opt_char}\.[^\x00]+?)(?=[A-D]\.|\x00|\x14|\x15|$)', temp_options_text)
                    if opt_match:
                        options.append(opt_match.group(1).strip().replace('\x00', ''))
                        temp_options_text = temp_options_text[opt_match.end():]
                    else:
                        if len(options) < 3:
                            # log_debug(f"后备逻辑在查找 '{opt_char}.' 时中断。")
                            options = []
                        break

                if len(options) != 4:
                     # log_debug(f"后备逻辑失败，找到 {len(options)} 个选项。")
                     return None, [], None

            # 清理可能混入的不可见字符或二进制片段
            cleaned_options = []
            for opt in options:
                # 找到第一个非文本常见字符的位置
                end_pos = -1
                for i, char in enumerate(opt):
                    # 允许字母数字、中文、标点、空格、制表符、换行符
                    if not ('\u4e00' <= char <= '\u9fff' or 'a' <= char.lower() <= 'z' or '0' <= char <= '9' or char in '().,，。（） \t\n'):
                         if i > 2: # 至少保留 A. 等前缀
                            end_pos = i
                            break
                if end_pos != -1:
                    cleaned_options.append(opt[:end_pos].strip())
                else:
                    cleaned_options.append(opt.strip())
            options = cleaned_options

            # 再次检查选项数量
            if len(options) != 4:
                 # log_debug(f"清理后选项数量 ({len(options)}) 仍不为 4，跳过。")
                 return None, [], None


            # log_debug(f"[PARSE_OK] 步骤 3 成功: 找到选项: {options}")

            # 4. 提取答案数组
            ans_match = re.search(b'\x01([\x30\x31])\x01([\x30\x31])\x01([\x30\x31])\x01([\x30\x31])', data)

            if not ans_match:
                # log_debug("[PARSE_FAIL] 步骤 4 失败: 未找到答案数组")
                return None, [], None

            answers = [ans_match.group(i).decode() for i in range(1, 5)]
            # log_debug(f"[PARSE_OK] 步骤 4 成功: 找到答案数组: {answers}")

            # 处理多选题，实际无用，多选题未在请求体暴露答案
            # 5. 查找所有正确答案索引
            correct_indices = [i for i, val in enumerate(answers) if val == '1']

            if not correct_indices:
                # log_debug("[PARSE_FAIL] 步骤 5 失败: 答案数组中没有 '1'。")
                return None, [], None

            # 检查索引是否有效
            if any(index >= len(options) for index in correct_indices):
                # log_debug(f"[PARSE_FAIL] 步骤 5 失败: 找到的正确答案索引 {correct_indices} 包含超出选项列表长度 {len(options)} 的值")
                return None, [], None

            # 组合所有正确答案的文本
            correct_answer_texts = [options[i] for i in correct_indices]
            correct_answer = ", ".join(correct_answer_texts) # 用逗号分隔多个答案
            # log_debug(f"[PARSE_OK] 步骤 5 成功: 正确答案为 {correct_answer}")

            return question, options, correct_answer

        except Exception as e:
            print(f"[ERROR] 解析 ViewState 时发生严重异常: {e}")
            import traceback
            traceback.print_exc()
            return None, [], None

    def request(self, flow: http.HTTPFlow):
        if not flow.request.pretty_url.startswith(TARGET_URL):
            return

        if flow.request.method != "POST" or not flow.request.content:
            return

        # log_debug(f"拦截到目标 POST 请求: {flow.request.pretty_url}")

        try:
            form_data = flow.request.get_text()
            parsed_data = urllib.parse.parse_qs(form_data)
        except Exception as e:
            print(f"[ERROR] 解析 POST 表单数据失败: {e}")
            return

        if '__VIEWSTATE' not in parsed_data:
            # log_debug("请求中未找到 __VIEWSTATE 字段。")
            return

        viewstate_b64 = parsed_data['__VIEWSTATE'][0]

        # log_debug(f"找到 __VIEWSTATE (长度: {len(viewstate_b64)})。正在转储到 {DEBUG_LOG_FILE}...") # 已禁用
        # dump_to_file(viewstate_b64) # 已禁用

        try:
            decoded_bytes = base64.b64decode(viewstate_b64.encode('latin-1'))
        except Exception as e:
            print(f"[ERROR] Base64 解码 ViewState 失败: {e}")
            return

        question, options, correct_answer = self.parse_viewstate_bytes(decoded_bytes)

        if not question:
            # log_debug("在此 ViewState 中未找到完整题目信息，跳过。") # 已禁用
            # print("-" * 40) # 已禁用
            return

        # 只在题目是新的时候才打印和保存
        if question not in self.known_questions:
            tid = flow.request.query.get('tid', 'unknown')

            print("=" * 40)
            print(f"✅ [嗅探成功] - 截获一道新题目")
            print(f"  TID: {tid}")
            print(f"  题目: {question}")
            for opt in options:
                print(f"    {opt}")
            print(f"  [!!] 正确答案: {correct_answer}")
            print("=" * 40)

            self.bank[question] = {
                'tid': tid,
                'question': question,
                'options': options,
                'correct_answer': correct_answer
            }
            self.known_questions.add(question) # 将新题目加入已知集合
            save_question_bank(self.bank)
            print(f"[INFO] 发现新题目！已保存到 {JSON_FILE}")
        # else: # 如果题目已知，则静默处理
            # log_debug(f"题目 '{question[:30]}...' 已存在于题库中，跳过打印。")

# Mitmproxy addon 入口
addons = [
    ViewStateSniffer()
]

if __name__ == "__main__":
    # 从 TARGET_URL 提取域名
    try:
        domain = urllib.parse.urlparse(TARGET_URL).hostname
        if not domain:
            raise ValueError("无法从 TARGET_URL 提取域名")
    except Exception as e:
        print(f"[CRITICAL] 无法解析 TARGET_URL: {e}")
        print("请确保 TARGET_URL 格式正确 (例如: http://example.com/path)")
        sys.exit(1)

    print("Mitmproxy ViewState嗅探脚本启动中...")
    print(f"请将你的设备代理指向 mitmproxy 运行的端口 (默认为 8080)")
    print(f"[INFO] 已设置控制台过滤器，将只显示来自 <{domain}> 的流量。")

    sys.argv.append('-s')
    sys.argv.append(__file__)
    sys.argv.append('-q')
    # 添加 mitmproxy 过滤器表达式，只显示目标域名的流量
    sys.argv.append(f"~d ^{domain}$") # 使用^和$确保精确匹配

    mitmdump()