tzzb 0.2.0__py3-none-any.whl

tzzb/__init__.py

@@ -0,0 +1,18 @@
+"""tzzb — Python SDK, CLI, and MCP Server for Tonghuashun Investment Ledger (同花顺投资账本)."""
+
+from .client import TzzbClient, TzzbError
+from .config import TzzbConfig, config
+from .models import Account, AssetTrendPoint, LoginStatus, Position, TradeRecord
+
+__all__ = [
+    "TzzbClient",
+    "TzzbError",
+    "TzzbConfig",
+    "config",
+    "Account",
+    "AssetTrendPoint",
+    "LoginStatus",
+    "Position",
+    "TradeRecord",
+]
+__version__ = "0.1.0"

tzzb/auth.py

@@ -0,0 +1,474 @@
+"""Authentication and Chrome DevTools Protocol (CDP) bridge for tzzb.
+
+Provides session management, Chrome process management, and CDP request execution
+to bypass Tonghuashun dynamic anti-bot signatures (hexin-v).
+"""
+
+from __future__ import annotations
+
+import json
+import socket
+import subprocess
+import threading
+import time
+import urllib.request
+from pathlib import Path
+from typing import Any
+
+import websocket
+
+from .models import LoginStatus
+
+COOKIE_FILE = Path.home() / ".tzzb_cookies.json"
+REQUIRED_COOKIES = {"userid", "ticket", "user"}
+TZZB_URL = "https://tzzb.10jqka.com.cn/pc/index.html#/unlogin"
+CDP_PORT = 9222
+CHROME_PROFILE = Path.home() / ".tzzb_chrome_profile"
+
+_ws: websocket.WebSocket | None = None
+_msg_id: int = 0
+_cdp_lock = threading.Lock()
+
+
+def check_port(port: int = CDP_PORT) -> bool:
+    """Check if the given port is open."""
+    try:
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        sock.settimeout(1)
+        result = sock.connect_ex(("127.0.0.1", port))
+        sock.close()
+        return result == 0
+    except Exception:
+        return False
+
+
+def find_chrome() -> str | None:
+    """Find system Chrome/Chromium executable."""
+    candidates = [
+        "/Applications/Google Chrome.app/Contents/MacOS/Google Chrome",
+        "/Applications/Chromium.app/Contents/MacOS/Chromium",
+        r"C:\Program Files\Google\Chrome\Application\chrome.exe",
+        r"C:\Program Files (x86)\Google\Chrome\Application\chrome.exe",
+        str(Path.home() / "AppData/Local/Google/Chrome/Application/chrome.exe"),
+        "/usr/bin/google-chrome",
+        "/usr/bin/chromium",
+        "/usr/bin/chromium-browser",
+    ]
+    for path in candidates:
+        if Path(path).exists():
+            return path
+
+    for cmd in [["which", "google-chrome"], ["which", "chromium"], ["which", "chrome"], ["where", "chrome"]]:
+        try:
+            result = subprocess.run(cmd, capture_output=True, text=True, timeout=5)
+            if result.returncode == 0 and result.stdout.strip():
+                return result.stdout.strip().split("\n")[0]
+        except Exception:
+            pass
+    return None
+
+
+def launch_chrome() -> bool:
+    """Launch isolated Chrome debug instance with visible GUI window."""
+    import sys
+    chrome_path = find_chrome()
+    if not chrome_path:
+        return False
+
+    CHROME_PROFILE.mkdir(parents=True, exist_ok=True)
+
+    if sys.platform == "darwin":
+        subprocess.Popen(
+            [
+                "open",
+                "-n",
+                "-a",
+                "Google Chrome",
+                "--args",
+                f"--remote-debugging-port={CDP_PORT}",
+                f"--user-data-dir={CHROME_PROFILE}",
+                "--remote-allow-origins=*",
+                "--no-first-run",
+                "--no-default-browser-check",
+                TZZB_URL,
+            ]
+        )
+    else:
+        creationflags = 0
+        if hasattr(subprocess, "DETACHED_PROCESS"):
+            creationflags |= subprocess.DETACHED_PROCESS
+        if hasattr(subprocess, "CREATE_NO_WINDOW"):
+            creationflags |= subprocess.CREATE_NO_WINDOW
+
+        subprocess.Popen(
+            [
+                chrome_path,
+                f"--remote-debugging-port={CDP_PORT}",
+                f"--user-data-dir={CHROME_PROFILE}",
+                "--remote-allow-origins=*",
+                "--no-first-run",
+                "--no-default-browser-check",
+                TZZB_URL,
+            ],
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+            creationflags=creationflags,
+        )
+
+    for _ in range(30):
+        time.sleep(1)
+        if check_port(CDP_PORT):
+            return True
+    return False
+
+
+def reset_connection() -> None:
+    """Close existing WebSocket connection and reset state."""
+    global _ws
+    if _ws:
+        try:
+            _ws.close()
+        except Exception:
+            pass
+        _ws = None
+
+
+def get_page_ws_url() -> str:
+    """Retrieve debugger WebSocket URL for an active page target."""
+    try:
+        resp = urllib.request.urlopen(f"http://127.0.0.1:{CDP_PORT}/json")
+        targets = json.loads(resp.read())
+    except Exception:
+        targets = []
+
+    for t in targets:
+        if t.get("type") == "page":
+            return t["webSocketDebuggerUrl"]
+
+    # Try creating a new target page if none exist
+    try:
+        req = urllib.request.Request(f"http://127.0.0.1:{CDP_PORT}/json/new?{TZZB_URL}", method="PUT")
+        resp = urllib.request.urlopen(req)
+        new_target = json.loads(resp.read())
+        if "webSocketDebuggerUrl" in new_target:
+            return new_target["webSocketDebuggerUrl"]
+    except Exception:
+        pass
+
+    raise RuntimeError("No available page targets in Chrome CDP")
+
+
+def wait_for_page_load(timeout: int = 15) -> None:
+    """Wait until page state is ready for requests."""
+    global _msg_id
+
+    _msg_id += 1
+    nav_msg = {
+        "id": _msg_id,
+        "method": "Page.navigate",
+        "params": {"url": "https://tzzb.10jqka.com.cn/pc/index.html"},
+    }
+    _ws.send(json.dumps(nav_msg))
+
+    try:
+        _ws.settimeout(5)
+        while True:
+            raw = _ws.recv()
+            parsed = json.loads(raw)
+            if parsed.get("id") == _msg_id:
+                break
+    except Exception:
+        pass
+
+    check_expr = "(function(){return document.readyState==='complete'?1:0})()"
+    deadline = time.time() + timeout
+    while time.time() < deadline:
+        _msg_id += 1
+        eval_msg = {
+            "id": _msg_id,
+            "method": "Runtime.evaluate",
+            "params": {
+                "expression": check_expr,
+                "returnByValue": True,
+                "awaitPromise": False,
+            },
+        }
+        _ws.send(json.dumps(eval_msg))
+
+        try:
+            _ws.settimeout(3)
+            while True:
+                raw = _ws.recv()
+                parsed = json.loads(raw)
+                if parsed.get("id") == _msg_id:
+                    value = parsed.get("result", {}).get("result", {}).get("value", 0)
+                    if value == 1:
+                        time.sleep(1)
+                        return
+                    break
+        except websocket.WebSocketTimeoutException:
+            pass
+        time.sleep(0.5)
+
+    time.sleep(2)
+
+
+def cdp_connect() -> None:
+    """Connect to Chrome CDP target."""
+    global _ws, _msg_id
+
+    reset_connection()
+    ws_url = get_page_ws_url()
+    _ws = websocket.create_connection(ws_url, timeout=10)
+    _msg_id = 0
+
+    cdp_send("Page.enable")
+    wait_for_page_load(timeout=15)
+
+
+def cdp_send(method: str, params: dict[str, Any] | None = None) -> dict[str, Any]:
+    """Send CDP command and wait for response."""
+    global _msg_id
+
+    _msg_id += 1
+    msg = {"id": _msg_id, "method": method, "params": params or {}}
+    _ws.send(json.dumps(msg))
+
+    while True:
+        raw = _ws.recv()
+        parsed = json.loads(raw)
+        if parsed.get("id") == _msg_id:
+            return parsed
+
+
+def ensure_connection() -> None:
+    """Ensure CDP connection is established and healthy."""
+    global _ws
+
+    if _ws is not None:
+        try:
+            _ws.ping()
+            return
+        except Exception:
+            reset_connection()
+
+    if not check_port(CDP_PORT):
+        if not launch_chrome():
+            raise RuntimeError("Failed to launch or connect to Chrome on CDP port 9222")
+
+    cdp_connect()
+
+
+def cdp_send_safe(method: str, params: dict[str, Any] | None = None) -> dict[str, Any]:
+    """Thread-safe CDP call with automatic reconnection."""
+    with _cdp_lock:
+        ensure_connection()
+        try:
+            return cdp_send(method, params)
+        except (websocket.WebSocketException, ConnectionError, OSError):
+            reset_connection()
+            ensure_connection()
+            try:
+                return cdp_send(method, params)
+            except Exception:
+                reset_connection()
+                raise
+
+
+def cdp_evaluate(expression: str) -> dict[str, Any]:
+    """Execute JavaScript in browser via Runtime.evaluate."""
+    return cdp_send_safe(
+        "Runtime.evaluate",
+        {
+            "expression": expression,
+            "returnByValue": True,
+            "awaitPromise": True,
+        },
+    )
+
+
+def load_cookies() -> dict[str, str] | None:
+    """Load cached cookies from disk."""
+    if not COOKIE_FILE.exists():
+        return None
+    try:
+        data = json.loads(COOKIE_FILE.read_text(encoding="utf-8"))
+        if REQUIRED_COOKIES.issubset(data.keys()):
+            return {k: str(data[k]) for k in REQUIRED_COOKIES}
+    except Exception:
+        pass
+    return None
+
+
+def save_cookies(cookies: dict[str, str]) -> None:
+    """Persist cookies to disk."""
+    data = {
+        **cookies,
+        "extracted_at": time.strftime("%Y-%m-%dT%H:%M:%S"),
+        "source": "cdp",
+    }
+    COOKIE_FILE.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
+
+
+def clear_cookies() -> None:
+    """Clear cached cookies."""
+    COOKIE_FILE.unlink(missing_ok=True)
+
+
+def cdp_get_cookies() -> dict[str, str]:
+    """Extract cookies from browser via Network.getCookies."""
+    result = cdp_send_safe("Network.getCookies", {"urls": ["https://tzzb.10jqka.com.cn"]})
+    cookies = result.get("result", {}).get("cookies", [])
+    extracted: dict[str, str] = {}
+    for c in cookies:
+        if "10jqka.com.cn" in c.get("domain", "") and c["name"] in REQUIRED_COOKIES:
+            extracted[c["name"]] = c["value"]
+    return extracted if REQUIRED_COOKIES.issubset(extracted.keys()) else {}
+
+
+def get_login_status() -> LoginStatus:
+    """Get current authentication status."""
+    cookies = load_cookies()
+    if cookies:
+        try:
+            raw = json.loads(COOKIE_FILE.read_text(encoding="utf-8"))
+            extracted_at = raw.get("extracted_at", "")
+        except Exception:
+            extracted_at = ""
+        return LoginStatus(
+            logged_in=True,
+            userid=cookies.get("userid", ""),
+            cookie_count=len(cookies),
+            extracted_at=extracted_at,
+            message="Authenticated",
+        )
+    return LoginStatus(logged_in=False, message="Not authenticated")
+
+
+def logout() -> dict[str, Any]:
+    """Log out by clearing both file and browser cookies."""
+    clear_cookies()
+    try:
+        if check_port(CDP_PORT):
+            cdp_send_safe("Network.clearBrowserCookies")
+            cdp_send_safe("Page.navigate", {"url": TZZB_URL})
+    except Exception:
+        pass
+    return {"status": "success", "message": "Logged out successfully"}
+
+
+def focus_chrome() -> None:
+    """Bring Chrome window to the foreground on macOS."""
+    import sys
+    if sys.platform == "darwin":
+        try:
+            subprocess.run(
+                ["osascript", "-e", 'tell application "Google Chrome" to activate'],
+                capture_output=True,
+                timeout=2,
+            )
+        except Exception:
+            pass
+
+
+def ensure_qr_login_active(timeout: float = 15.0) -> bool:
+    """Ensure the login dialog is open and switched to QR code scan mode."""
+    deadline = time.time() + timeout
+    while time.time() < deadline:
+        # 1. Trigger modal on main page if not open
+        trigger_modal_js = """
+        (() => {
+            const popup = document.getElementById('ths_iframe_login_popup');
+            if (!popup || popup.style.display === 'none') {
+                const btn = document.querySelector('.UnloginBar_login') || document.querySelector('.Unlogin_top_box_btn');
+                if (btn) btn.click();
+                return true;
+            }
+            return false;
+        })()
+        """
+        try:
+            cdp_evaluate(trigger_modal_js)
+        except Exception:
+            pass
+
+        # 2. Check child iframe and ensure qrcode_pannel is active
+        try:
+            ft = cdp_send_safe("Page.getFrameTree")
+            child_frames = ft.get("result", {}).get("frameTree", {}).get("childFrames", [])
+            for cf in child_frames:
+                frame_id = cf.get("frame", {}).get("id")
+                if not frame_id:
+                    continue
+                world = cdp_send_safe("Page.createIsolatedWorld", {"frameId": frame_id, "worldName": "tzzb_login"})
+                ctx_id = world.get("result", {}).get("executionContextId")
+                if ctx_id:
+                    res = cdp_send_safe(
+                        "Runtime.evaluate",
+                        {
+                            "contextId": ctx_id,
+                            "expression": """
+                            (() => {
+                                const qrPannel = document.getElementById('qrcode_pannel');
+                                if (qrPannel && window.getComputedStyle(qrPannel).display === 'block') {
+                                    return 'ready';
+                                }
+                                const toQr = document.getElementById('to_qrcode_login');
+                                if (toQr) {
+                                    toQr.click();
+                                    return 'clicked';
+                                }
+                                return 'waiting';
+                            })()
+                            """,
+                            "returnByValue": True,
+                        },
+                    )
+                    status = res.get("result", {}).get("result", {}).get("value")
+                    if status == "ready":
+                        return True
+        except Exception:
+            pass
+        time.sleep(0.5)
+    return False
+
+
+def login(timeout: int = 120, force: bool = False) -> dict[str, Any]:
+    """Interactive login flow via Chrome CDP."""
+    try:
+        if not check_port(CDP_PORT):
+            if not launch_chrome():
+                return {"status": "error", "message": "Failed to launch Chrome. Please launch Chrome manually."}
+
+        if force:
+            logout()
+
+        existing = cdp_get_cookies()
+        if existing and not force:
+            save_cookies(existing)
+            return {
+                "status": "success",
+                "message": f"Cookie valid, userid={existing['userid'][:6]}***",
+                "cookies": {k: v[:6] + "***" for k, v in existing.items()},
+            }
+
+        cdp_send_safe("Page.navigate", {"url": TZZB_URL})
+        focus_chrome()
+        time.sleep(1)
+        ensure_qr_login_active(timeout=15.0)
+
+        start_time = time.time()
+        while time.time() - start_time < timeout:
+            time.sleep(1)
+            cookies = cdp_get_cookies()
+            if cookies:
+                save_cookies(cookies)
+                return {
+                    "status": "success",
+                    "message": f"Login successful, userid={cookies['userid'][:6]}***",
+                    "cookies": {k: v[:6] + "***" for k, v in cookies.items()},
+                }
+
+        return {"status": "timeout", "message": f"Login timed out after {timeout} seconds"}
+    except Exception as e:
+        return {"status": "error", "message": str(e)}