coordinate/plugins/bel_plugin.py

import time
import json
import random
import re
import os
import uuid
import shutil
import base64
import socket
import requests
from datetime import datetime
from typing import List, Dict, Optional, Any, Callable
from urllib.parse import urljoin, urlparse, urlencode

# DrissionPage 核心
from DrissionPage import ChromiumPage, ChromiumOptions


from vs_plg import IVSPlg
from vs_types import VSPlgConfig, AppointmentType, VSQueryResult, VSBookResult, AvailabilityStatus, TimeSlot, DateAvailability, NotFoundError, PermissionDeniedError, RateLimiteddError, SessionExpiredOrInvalidError, BizLogicError 
from toolkit.vs_cloud_api import VSCloudApi
from toolkit.proxy_tunnel import ProxyTunnel


class BrowserResponse:
    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

def to_yyyymmdd(data_str: str, date_str_format: str, target_format: str="%Y-%m-%d"):
    dt = datetime.strptime(data_str, date_str_format)
    return dt.strftime("%Y-%m-%d")

def get_alias_email(email: str, new_domain: str = "gmail-app.com") -> str:
    if "@" not in email: raise ValueError(f"Invalid email: {email}")
    local_part, _ = email.rsplit("@", 1)
    return f"{local_part}@{new_domain}"

class BelPlugin(IVSPlg):
    """
    Belgium (https://visaonweb.diplomatie.be/en) 签证预约插件 (Browser + Tunnel Mode)
    """

    def __init__(self, group_id: str):
        self.group_id = group_id
        self.config: Optional[VSPlgConfig] = None
        self.free_config: Dict[str, Any] = {}
        self.logger = 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.applicants = []
        
        self.tunnel = None # 代理隧道
        self.is_healthy = True
        self.session_create_time: float = 0

    def get_group_id(self) -> str:
        return self.group_id
    
    def set_log(self, logger: Callable[[str], None]) -> None:
        self.logger = logger
        
    def _log(self, message): 
        if self.logger:
            self.logger(f'[PolPlugin] [{self.group_id}] {message}')
        else:
            print(f'[PolPlugin] [{self.group_id}] {message}')
        
    def set_config(self, config: VSPlgConfig):
        self.config = config
        self.free_config = config.free_config or {}
            
    def keep_alive(self):
        pass

    def health_check(self) -> bool:
        if not self.is_healthy:
            return False
        if not self.page:
            return False
        try:
            if not self.page.run_js("return 1;"):
                return False
        except:
            return False
        if self.config.session_max_life > 0:
            if time.time() - self.session_create_time > self.config.session_max_life:
                self._log("Session expired.")
                return False
        return True

    def create_session(self):
        """
        创建会话：启动浏览器 -> 代理隧道 -> 提取 Captcha -> 本地识别 -> 提交 -> 获取 Context
        """
        self._log(f"Initializing Session (ID: {self.instance_id})...")
        co = ChromiumOptions()
        def get_free_port():
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                s.bind(('', 0)); return s.getsockname()[1]
        co.set_local_port(get_free_port())
        
        co.set_user_data_path(self.user_data_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 Tunnel for {p.ip}...")
                self.tunnel = ProxyTunnel(p.ip, p.port, p.username, p.password)
                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.scheme}://{p.ip}:{p.port}"
                co.set_argument(f'--proxy-server={proxy_str}')
        else:
            self._log("[WARN] No proxy configured!")

        co.headless(False)
        co.set_argument('--no-sandbox')
        co.set_argument('--disable-gpu')
        co.set_argument('--disable-dev-shm-usage') 
        co.set_argument('--window-size=1920,1080')
        co.set_argument('--disable-blink-features=AutomationControlled')
        co.set_argument('--ignore-certificate-errors')

        try:
            self.page = ChromiumPage(co)
            url_home = "https://visaonweb.diplomatie.be/Account/Login?ReturnUrl=%2Fen"
            self._log(f"Navigating to {url_home}")
            self.page.get(url_home)
            self.page.wait.doc_loaded()
            self.page.wait.ele_displayed('tag:button@@type=submit', timeout=10)
            self.page.ele('tag:input@@name=UserName').input(self.config.account.username)
            self.page.ele('tag:input@@name=Password').input(self.config.account.password)
            self.page.ele('tag:button@@type=submit').click()
            self.page.listen.start('/VisaApplication/MyList')
            self.page.get('https://visaonweb.diplomatie.be/en/VisaApplication/IndexByUserId')
          
            packet = self.page.listen.wait(timeout=10)
            if not packet:
                raise BizLogicError(message='Get userList failed')
            self.page.listen.stop()
            resp_body = packet.response.body
            for dat in resp_body.get('data', []):
                if dat.get('St') == 'Submitted':
                    applicant = {
                        "id": dat.get('Id'),
                        "first_name": dat.get('FName'),
                        "last_name": dat.get('LName'),
                        "vow_id": dat.get('VOWId'),
                        "os_id": dat.get('OSId'),
                        "sub_group_id": dat.get('SubGroupId'),
                        "vac": dat.get('Vac'),
                    }
                    self.applicants.append(applicant)
                
            self.session_create_time = time.time()
            self._log("Session created successfully.")
        except Exception as e:
            self._log(f"Session Create Failed: {e}")
            time.sleep(3600)
            self.cleanup()
            raise e

    def query(self, apt_type: AppointmentType) -> VSQueryResult:
        res = VSQueryResult()
        res.success = False
                   
        applicant = random.choice(self.applicants)
        applicant_id = applicant.get('id')
        url = 'https://visaonweb.diplomatie.be/Common/GetEAppointmentUrl'
        params = {
            'id': applicant_id
        }
        headers = {
            "cache-control": "max-age=0",
            "if-modified-since": "0",
            "x-requested-with": "XMLHttpRequest",
        }
        resp = self._perform_request('GET', url, headers=headers, params=params)
        appointment_url = resp.json().get('url')
        
        tab = self.page.new_tab(appointment_url)
        # solve hcaptcha
        # sitekey = 5f64399c-14a8-415e-ad1a-7ebccdc4943a
        
        token = self.get_captcha_solution(self.free_config.get('capsolver_key'), '5f64399c-14a8-415e-ad1a-7ebccdc4943a', appointment_url)
        script = f'document.querySelector(\'textarea[name="h-captcha-response"]\').value="{token}";'
        tab.run_js(script)
        time.sleep(5)
        self._log(tab.html)
        available_dates = []
        if available_dates:
            res.success = True
            res.availability_status = AvailabilityStatus.Available
            earliest_date = available_dates[0]
            earliest_dt = datetime.strptime(earliest_date, "%Y-%m-%d")
            res.earliest_date = earliest_dt

            res.availability = [
                DateAvailability(
                    date=datetime.strptime(d, "%Y-%m-%d"),
                    times=[],
                )
                for d in available_dates
            ]
        else:
            res.success = False
            res.availability_status = AvailabilityStatus.NoneAvailable
            res.availability = []
        return res
    
    def get_captcha_solution(self, api_key, site_key, site_url):
        payload = {
            "clientKey": api_key,
            "task": {
                "type": "HCaptchaTaskProxyLess",
                "websiteKey": site_key,
                "websiteURL": site_url
            }
        }
        
        with requests.Session() as session:
            res = session.post("https://api.capsolver.com/createTask", json=payload)
            task_id = res.json().get("taskId")
            if not task_id:
                raise Exception("Failed to create task:", res.text)
            
            print("Waiting for CAPTCHA solution...")
            while True:
                time.sleep(5)  # Wait for 5 seconds before checking the result
                res = session.post("https://api.capsolver.com/getTaskResult", json={"clientKey": api_key, "taskId": task_id})
                result = res.json()
                if result.get("status") == "ready":
                    return result["solution"]["gRecaptchaResponse"]
                if result.get("status") == "failed":
                    raise Exception("CAPTCHA solve failed:", res.text)

    def book(self, slot_info: VSQueryResult, user_inputs: Dict) -> VSBookResult:
        res = VSBookResult()
        return res
    
    def _perform_request(self, method, url, headers=None, data=None, json_data=None, params=None, retry_count=0):
        if not self.page:
            raise BizLogicError("Browser not init")
        
        req_url = url
        if params:
            sep = '&' if '?' in req_url else '?'
            req_url += sep + urlencode(params)
            
        fetch_opts = { "method": method.upper(), "headers": headers or {}, "credentials": "include" }
        
        if json_data:
            fetch_opts['body'] = json.dumps(json_data)
            fetch_opts['headers']['Content-Type'] = 'application/json'
        elif data:
            if isinstance(data, dict):
                fetch_opts['body'] = urlencode(data)
                fetch_opts['headers']['Content-Type'] = 'application/x-www-form-urlencoded; charset=UTF-8'
            else:
                fetch_opts['body'] = data

        js = f"""
        return fetch("{req_url}", {json.dumps(fetch_opts)})
        .then(async r => {{
            const h = {{}}; r.headers.forEach((v, k) => h[k] = v);
            return {{ status: r.status, body: await r.text(), headers: h, url: r.url }};
        }}).catch(e => {{ return {{ status: 0, body: e.toString() }}; }});
        """
        
        resp = BrowserResponse(self.page.run_js(js, timeout=60))
        
        if resp.status_code == 200:
            return resp
        elif resp.status_code == 403:
            if "Just a moment" in resp.text and retry_count < 2:
                self._log("Cloudflare 403. Refreshing...")
                if self._refresh_firewall_session():
                    return self._perform_request(method, url, headers, data, json_data, params, retry_count+1)
            raise PermissionDeniedError(f"HTTP 403: {resp.text[:100]}")
        elif resp.status_code == 429:
            self.is_healthy = False
            raise RateLimiteddError()
        elif resp.status_code in [401, 419]:
            self.is_healthy = False
            raise SessionExpiredOrInvalidError()
        else:
            raise BizLogicError(f"HTTP {resp.status_code}: {resp.text[:100]}")

    def _filter_dates(self, dates, start, end):
        if not start or not end: return dates
        valid = []
        s = datetime.strptime(start[:10], "%Y-%m-%d")
        e = datetime.strptime(end[:10], "%Y-%m-%d")
        for d in dates:
            c = datetime.strptime(d, "%Y-%m-%d")
            if s <= c <= e: valid.append(d)
        random.shuffle(valid)
        return valid

    def cleanup(self):
        if self.page:
            try: self.page.quit()
            except: pass
            self.page = None
        if os.path.exists(self.root_workspace):
            for _ in range(3):
                try: time.sleep(0.2); shutil.rmtree(self.root_workspace, ignore_errors=True); break
                except: time.sleep(0.5)
        if self.tunnel:
            try: self.tunnel.stop()
            except: pass
            self.tunnel = None
        
    def __del__(self):
        self.cleanup()