live-hub-py/core/sms_provider.py

"""接码平台查询与验证码提取。"""

from __future__ import annotations

import json
import re
from dataclasses import dataclass
from datetime import UTC, datetime
from urllib.parse import urlparse

import requests

CODE_PATTERN = re.compile(
    r"(?:verification\s+code|验证码|校验码|动态码|安全码)\D{0,20}(\d{4,8})",
    re.IGNORECASE,
)
FALLBACK_CODE_PATTERN = re.compile(r"(?<!\d)(\d{4,8})(?!\d)")


@dataclass
class SmsLine:
    """手机号和短信查询地址。"""

    phone: str
    url: str
    provider: str
    raw: str


@dataclass
class SmsPollResult:
    """短信查询结果。"""

    status: str
    code: str = ""
    message: str = ""
    raw: str = ""

    @property
    def has_code(self) -> bool:
        return self.status == "code" and bool(self.code)


def parse_sms_lines(text: str) -> list[SmsLine]:
    """解析手机号池，兼容手机号----查询URL、账号----密码----手机号----查询URL。"""
    rows: list[SmsLine] = []
    for raw_line in (text or "").splitlines():
        line = raw_line.strip()
        if not line or line.startswith("#"):
            continue
        parts = [part.strip() for part in line.split("----")]
        if len(parts) >= 4 and _looks_like_url("----".join(parts[3:])):
            phone, url = parts[2], "----".join(parts[3:]).strip()
        elif len(parts) >= 2 and _looks_like_url("----".join(parts[1:])):
            phone, url = parts[0], "----".join(parts[1:]).strip()
        else:
            continue
        if not phone or not url:
            continue
        rows.append(
            SmsLine(
                phone=phone,
                url=url,
                provider=detect_provider(url),
                raw=line,
            )
        )
    return rows


def _looks_like_url(value: str) -> bool:
    """判断文本是否像 URL。"""
    parsed = urlparse(str(value or "").strip())
    return parsed.scheme in {"http", "https"} and bool(parsed.netloc)


def detect_provider(url: str) -> str:
    """根据 URL 判断接码平台类型。"""
    text = (url or "").lower()
    if "sms8.net" in text:
        return "sms8"
    return "pipe"


def extract_sms_code(text: str) -> str:
    """从短信内容中提取验证码。"""
    body = str(text or "")
    match = CODE_PATTERN.search(body)
    if match:
        return match.group(1)
    match = FALLBACK_CODE_PATTERN.search(body)
    return match.group(1) if match else ""


def _parse_sms8_time(value: str) -> datetime | None:
    text = str(value or "").strip()
    if not text:
        return None
    for fmt in ("%Y-%m-%d %H:%M:%S", "%Y/%m/%d %H:%M:%S"):
        try:
            return datetime.strptime(text, fmt).replace(tzinfo=UTC)
        except ValueError:
            continue
    return None


class SmsProviderClient:
    """固定手机号池的短信查询客户端。"""

    def __init__(self, timeout: tuple[float, float] = (8, 20)):
        self.timeout = timeout
        self.session = requests.Session()
        self.session.trust_env = False

    def poll(self, item: SmsLine, after_time: datetime | None = None) -> SmsPollResult:
        """查询一次短信。"""
        try:
            response = self.session.get(item.url, timeout=self.timeout)
            response.raise_for_status()
        except requests.RequestException as exc:
            return SmsPollResult(status="error", message=f"查询短信失败: {exc}")

        body = response.text.strip()
        if item.provider == "sms8":
            return self._parse_sms8(body, after_time=after_time)
        return self._parse_pipe(body)

    def _parse_pipe(self, body: str) -> SmsPollResult:
        if not body:
            return SmsPollResult(status="waiting", message="短信为空", raw=body)
        low = body.lower()
        if low.startswith("no|"):
            return SmsPollResult(status="waiting", message=body, raw=body)
        if low.startswith("yes|"):
            code = extract_sms_code(body)
            if code:
                return SmsPollResult(
                    status="code", code=code, message="收到验证码", raw=body
                )
            return SmsPollResult(
                status="error", message=f"已收到短信但未识别验证码: {body}", raw=body
            )
        code = extract_sms_code(body)
        if code:
            return SmsPollResult(
                status="code", code=code, message="收到验证码", raw=body
            )
        return SmsPollResult(status="waiting", message=body, raw=body)

    def _parse_sms8(
        self, body: str, after_time: datetime | None = None
    ) -> SmsPollResult:
        try:
            payload = json.loads(body)
        except json.JSONDecodeError:
            code = extract_sms_code(body)
            if code:
                return SmsPollResult(
                    status="code", code=code, message="收到验证码", raw=body
                )
            return SmsPollResult(
                status="error", message=f"短信平台返回非 JSON: {body[:120]}", raw=body
            )

        data = payload.get("data") if isinstance(payload.get("data"), dict) else {}
        msg = str(payload.get("msg") or "")
        code_text = str(data.get("code") or "")
        code_time = _parse_sms8_time(str(data.get("code_time") or ""))
        if after_time and code_time and code_time < after_time.replace(tzinfo=None):
            return SmsPollResult(
                status="waiting", message="验证码时间早于本次发码", raw=body
            )

        code = extract_sms_code(code_text)
        if int(payload.get("code") or 0) == 1 and code:
            return SmsPollResult(
                status="code", code=code, message="收到验证码", raw=body
            )
        if code:
            return SmsPollResult(
                status="code", code=code, message="收到验证码", raw=body
            )
        return SmsPollResult(status="waiting", message=msg or "暂无验证码", raw=body)