coordinate/plugins/usa_plugin.py

import time
import json
import random
import re
import os
import uuid
import shutil
import socket
from datetime import date, datetime
from typing import List, Dict, Optional, Any, Callable
from urllib.parse import urljoin, urlparse, urlencode, parse_qs
from concurrent.futures import ThreadPoolExecutor
from DrissionPage import ChromiumPage, ChromiumOptions

import configure
from vs_plg import IVSPlg
from utils.cloudflare_bypass_for_scraping import CloudflareBypasser
from toolkit.mihomo_tunnel import MihomoTunnel
from utils.mouse import HumanMouse
from utils.keyboard import HumanKeyboard
from utils.fingerprint_utils import FingerprintGenerator
from vs_types import VSPlgConfig, AppointmentType, VSQueryResult, VSBookResult, AvailabilityStatus, TimeSlot, DateAvailability, NotFoundError, PermissionDeniedError, RateLimiteddError, SessionExpiredOrInvalidError, BizLogicError 


class BrowserResponse:
    """模拟 requests.Response"""
    def __init__(self, result_dict):
        result_dict = result_dict or {}
        self.status_code = result_dict.get('status', 0)
        self.text = result_dict.get('body', '')
        self.headers = result_dict.get('headers', {})
        self.url = result_dict.get('url', '')
        self._json = None

    def json(self):
        if self._json is None:
            if not self.text:
                return {}
            try:
                self._json = json.loads(self.text)
            except:
                self._json = {}
        return self._json


class UsaPlugin(IVSPlg):
    
    LOCATIONS = {
        "SHANGHAI": {"name": "SHANGHAI", "id": "096bf614-b0db-ec11-a7b4-001dd80234f6"},
        "WUHAN": {"name": "WUHAN", "id": "7b6af614-b0db-ec11-a7b4-001dd80234f6"},
        "SHENYANG": {"name": "SHENYANG", "id": "0f6bf614-b0db-ec11-a7b4-001dd80234f6"},
    }

    def __init__(self, group_id: str):
        self.group_id = group_id
        self.config: Optional[VSPlgConfig] = None
        self.free_config: Dict[str, Any] = {}
        self.is_healthy = True
        self.logger = None
        
        self.mouse = None
        self.keyboard = None
        self.page: Optional[ChromiumPage] = None
        
        self.instance_id = uuid.uuid4().hex[:8]
        self.root_workspace = os.path.abspath(os.path.join("data/temp_browser_data", f"{self.group_id}.{self.instance_id}"))
        self.user_data_path = os.path.join(self.root_workspace, "user_data")
        
        if not os.path.exists(self.root_workspace):
            os.makedirs(self.root_workspace)     

        self.tunnel = None
        self.session_create_time: float = 0

    def _log(self, message):
        if self.logger:
            self.logger(f'[UsaPlugin] [{self.group_id}] [{self.instance_id}] {message}')
        else:
            print(f'[UsaPlugin] [{self.group_id}] [{self.instance_id}] {message}')

    def _random_sleep(self, min_sec=30, max_sec=60):
        """
        核心防限速控制：模拟人类阅读和操作的长时间停顿，确保网络请求间隔在30-60秒
        """
        sleep_time = random.uniform(min_sec, max_sec)
        self._log(f"Anti-Rate-Limit: Sleeping for {sleep_time:.2f} seconds...")
        time.sleep(sleep_time)

    def _get_free_port(self):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.bind(('', 0))
            return s.getsockname()[1]

    def set_config(self, config: VSPlgConfig):
        """设置 API 的配置信息"""
        self.config = config
        self.free_config = config.free_config or {}

    def set_log(self, logger: Callable[[str], None]) -> None:
        """设置日志输出工具"""
        self.logger = logger
        
    def keep_alive(self):
        pass

    def health_check(self) -> bool:
        if not self.is_healthy:
            return False
        if self.page is None:
            return False
        try:
            if not self.page.run_js("return 1;"):
                return False
        except:
            return False
            
        if self.config.session_max_life > 0:
            current_time = time.time()
            elapsed_time = current_time - self.session_create_time
            if elapsed_time > self.config.session_max_life:
                self._log(f"Session expired.")
                return False
        return True
    
    def _save_screenshot(self, name_prefix):
        try:
            timestamp = int(time.time())
            filename = f"{self.instance_id}_{name_prefix}_{timestamp}.jpg"
            save_path = os.path.join("data", filename)
            os.makedirs("data", exist_ok=True)
            self.page.get_screenshot(path=save_path, full_page=False)
            self._log(f"Screenshot saved to {save_path}")
        except Exception as e:
            self._log(f"Failed to save screenshot: {e}")

    def create_session(self) -> None:
        """创建一个新的会话 (包含初始化浏览器、过CF验证和执行登录)"""
        
        self._log(f"Initializing Session (ID: {self.instance_id})...")

        def get_free_port():
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                s.bind(('', 0))
                return s.getsockname()[1]
        
        co = ChromiumOptions()
        debug_port = get_free_port()
        self._log(f"Assigned Debug Port: {debug_port}")
        self._log(f"Account id={self.config.account.id}, proxy id={self.config.proxy.id}")
        
        co.set_local_port(debug_port)       
        co.set_user_data_path(self.user_data_path)
        
        chrome_path = configure.CHROME_PATH
        if not chrome_path:
            chrome_path = os.getenv("CHROME_BIN")
        if chrome_path and os.path.exists(chrome_path):
            co.set_paths(browser_path=chrome_path)
        
        if self.config.proxy and self.config.proxy.ip:
            p = self.config.proxy
            if p.username and p.password:
                self._log(f"Starting Proxy Tunnel for {p.ip}...")
                exit_node = {
                    "name": "ExitNode",
                    "type": p.proto,
                    "server": p.ip,
                    "port": p.port,
                    "username": p.username,
                    "password": p.password
                }
                relay_node = None
                if configure.MIHOMO_RELAY_NODES:
                    relay_node = random.choice(configure.MIHOMO_RELAY_NODES)
                mihomo_path = configure.MIHOMO_BIN_PATH
                if not mihomo_path:
                    mihomo_path = os.getenv("MIHOMO_BIN")
                if not mihomo_path:
                    raise BizLogicError(message='Mihomo path is null, You need set mihomo bin path in configure or os env')
                self.tunnel = MihomoTunnel(mihomo_path, exit_node=exit_node, relay_node=relay_node)
                local_proxy = self.tunnel.start()
                self._log(f"Tunnel started at {local_proxy}")
                co.set_argument(f'--proxy-server={local_proxy}')
            else:
                proxy_str = f"{p.proto}://{p.ip}:{p.port}"
                co.set_argument(f'--proxy-server={proxy_str}')
        else:
            self._log("[WARN] No proxy configured!")

        specific_fp = FingerprintGenerator().generate(self.config.account.username)
        fp_seed = specific_fp.get("seed")
        fp_platform = specific_fp.get("platform")
        fp_brand = specific_fp.get("brand")
        self._log(f'browser fingerprint seed={fp_seed}')
        
        co.headless(False)
        co.set_argument('--no-sandbox')
        co.set_argument('--disable-dev-shm-usage')
        co.set_argument('--window-size=1920,1080')
        co.set_argument('--disable-blink-features=AutomationControlled')
        co.set_argument(f"--fingerprint={fp_seed}")
        co.set_argument(f"--fingerprint-platform={fp_platform}")
        co.set_argument(f"--fingerprint-brand={fp_brand}")
        
        self.page = ChromiumPage(co)
        # 获取基础 URL
        usa_url = self.free_config.get('usa_url', '') 
        self._log(f"Navigating: {usa_url}")
        self.page.get(usa_url)
        # 初始访问页面，等待长延时，防止过快触发后续操作
        self._random_sleep(30, 45)
        
        if 'Attention Required! | Cloudflare' in self.page.title and 'Sorry, you have been blocked' in self.page.html:
            self._log(f'Block by cloudflare, try refresh...')
            self.page.refresh()
            self._random_sleep(30, 45) # 刷新动作属于高危操作，加长延时
            self.page.wait.doc_loaded()
        
        cf_bypasser = CloudflareBypasser(self.page, log=self.config.debug)
        if not cf_bypasser.bypass(max_retry=6):
            raise BizLogicError("Cloudflare bypass timeout")
        # 绕过盾后，休眠一段时间再处理 waiting room
        self._random_sleep(15, 30) 
        cf_bypasser.handle_waiting_room()
        
        self._log("Init humanize tools...")
        self.mouse = HumanMouse(self.page, debug=False)
        self.keyboard = HumanKeyboard(self.page)
        
        viewport_width = self.page.rect.viewport_size[0]
        viewport_height = self.page.rect.viewport_size[1]
        init_x = random.randint(10, viewport_width - 10)
        init_y = random.randint(10, viewport_height - 10)
        self.mouse.move(init_x, init_y) 
        
        username = self.config.account.username
        password = self.config.account.password
        security = self.free_config.get('security', {})

        max_steps = 15 # 由于状态多，步数可以稍微调大一点
        stuck_counter = 0
        last_url = ""
        session_created = False
        has_submitted_login = False
        
        for step in range(max_steps):
            self.page.wait.doc_loaded()
            time.sleep(1) # 这个用于等待页面DOM渲染，保留短时，因为不是发请求
            
            current_url = self.page.url
            current_title = self.page.title.lower()
            current_html_content = self.page.html
            self._log(f"--- [Router Step {step+1}] Current URL: {current_url} ---")
            
            if current_url == last_url:
                stuck_counter += 1
            else:
                last_url = current_url
                stuck_counter = 0
            
            # --- [异常处理层] ---
            if stuck_counter >= 3:
                self._log("[WARN] Page stucked, try to refresh...")
                self.page.refresh()
                self._random_sleep(30, 60) # 刷新操作触发请求，长延时
                stuck_counter = 0
                continue
            
            server_error_indicators = ["502 Bad Gateway", "503 Service Temporarily Unavailable"]
            # 网络出现故障,直接重试
            if any(err in current_html_content for err in server_error_indicators):
                self._log(f"[WARN] Server network error, try to refresh (Step: {step})...")
                self.page.refresh()
                self._random_sleep(45, 60) # 遇到错误必须拉长延时防止被彻底拉黑
                continue
            
            cloudflare_blocked_indicators = [
                "Sorry, you have been blocked" in current_html_content,
                "You are being rate limited" in current_html_content,
                "Cloudflare Ray ID" in current_html_content
            ]
            if any(cloudflare_blocked_indicators):
                raise BizLogicError(message="Blocked by Cloudflare WAF. Need to change IP or browser fingerprint.")
            
            # 遇到五秒盾先绕盾
            if "just a moment" in current_title:
                cf_bypasser.bypass(max_retry=3)
                self._random_sleep(20, 40)
                continue
            
            if self.page.ele('#post_select', timeout=1):
                self._log("🎉 Successfully reached the Slot Search page (Target Page). Session created successfully!")                
                self.session_create_time = time.time()
                session_created = True
                break

            # 状态 2: 密保问题页面
            elif self.page.ele('xpath://input[starts-with(@id, "kba") and contains(@id, "_response")]', timeout=1):
                self._log("[State] Security question verification detected. Filling in answers...")
                answer_eles = self.page.eles('xpath://input[starts-with(@id, "kba") and contains(@id, "_response")]')
                
                for ans_ele in answer_eles:
                    ele_id = ans_ele.attr('id')
                    match = re.search(r'kba(\d+)_response', ele_id)
                    if match:
                        q_num = match.group(1)
                        config_key = f"{q_num}_quest"
                        
                        q_data = security.get(config_key)
                
                        ans_text = q_data.get('a')
                        ans_ele.input(ans_text)
                        self._log(f"-> Find input {ele_id}, successfully filled in the answer for question {q_num}.")
          

                self.page.ele('#continue').click()
                self._log("Security answers submitted. Waiting for redirection...")
                # 提交密保问题，触发POST请求，长延时
                self._random_sleep(30, 60) 
                continue

            # 状态 1: 登录页面
            elif self.page.ele('#signInName', timeout=1):
                self._log("[State] Login page detected. Submitting credentials...")
                
                username_input = self.page.ele('#signInName')
                username_input.clear()
                username_input.input(username)
                
                password_input = self.page.ele('#password')
                password_input.clear()
                password_input.input(password)
                
                self.page.ele('#continue').click()
                has_submitted_login = True
                self._log("Login form submitted. Waiting for the next step to load...")
                # 提交登录表单，触发POST请求，必须长延时
                self._random_sleep(30, 60)
                continue

            # 状态 3: 预约主页(控制台) -> 选择首签或改签
            elif self.page.ele('#atlas-sidebar', timeout=1):
                self._log("[State] At the main booking dashboard. Looking for navigation button...")
                
                reschedule_btn = self.page.ele('#reschedule_appointment', timeout=0.5)
                if reschedule_btn:
                    self._log("-> Detected [Reschedule Appointment]. Currently in rescheduling mode, clicking to proceed...")
                    reschedule_btn.click()
                    # 点击导航按钮，触发页面跳转GET请求，长延时
                    self._random_sleep(30, 50)
                else:
                    schedule_btn = self.page.ele('xpath://ul[@id="atlas-sidebar"]//a[text()="安排预约" or text()="New Appointment" or text()="Schedule Appointment"]', timeout=0.5)
                    if schedule_btn:
                        self._log("-> Detected [Schedule Appointment]. Currently in first-time booking mode, clicking to proceed...")
                        schedule_btn.click()
                        # 点击导航按钮，触发页面跳转GET请求，长延时
                        self._random_sleep(30, 50)
                    else:
                        self._log("-> [WARN] Sidebar found, but no 'Schedule' or 'Reschedule' button detected. The page may still be loading...")
                        time.sleep(2) # 仅等待DOM渲染，不发请求，保持短时
                continue

            else:
                self._log("[State] In unknown or transitional state. No matching UI elements found. Waiting for next polling cycle...")
                time.sleep(2) # 仅等待DOM渲染，不发请求，保持短时
                
        if not session_created:
            raise BizLogicError(f"Failed to reach appointment-booking after {max_steps} navigation steps. Stuck at: {self.page.url}")


    def query(self, apt_type: AppointmentType) -> VSQueryResult:
        """查询可用的签证预约信息"""
        self._log("Querying available slots...")
        res = VSQueryResult()
        res.success = False
        
        # 1. 刷新页面以获取最新数据
        self.page.refresh()
        # 【关键修改】：刷新操作触发页面重载请求，必须执行长睡眠 
        self._random_sleep(30, 60)
        
        current_url = self.page.url.lower()
        if 'auth' in current_url or 'login' in current_url:
            self.is_healthy = False
            raise SessionExpiredOrInvalidError()
        
        applicant = self.free_config.get('applicant')
        location_name = self.free_config.get('location')
        location_id = self.LOCATIONS.get(location_name.upper(), {}).get('id')

        # 2. 等待页面元素
        self.page.ele(f"xpath://label[text()='{applicant}']", timeout=60)
        post_select = self.page.ele('#po