app.polisplexity.tech/pxy_bots/router.py

# pxy_bots/router.py
import json
import logging
import time  # <-- add
from typing import Dict, Optional, Tuple
from urllib.parse import urlparse

logger = logging.getLogger(__name__)

# Try to use requests; fallback to urllib d
try:
    import requests  # type: ignore
    _HAS_REQUESTS = True
except Exception:
    import urllib.request  # type: ignore
    _HAS_REQUESTS = False


def _is_allowed(url: str, allowed_hosts: Optional[set]) -> Tuple[bool, str]:
    """Allowlist check for DB routes (per-Connection)."""
    try:
        p = urlparse(url)
        if p.scheme not in {"http", "https"}:
            return False, "bad_scheme"
        host = (p.hostname or "").lower()
        return (host in (allowed_hosts or set())), f"host={host}"
    except Exception as e:
        return False, f"invalid_url:{e}"


def _compose_url(base: str, path: str) -> str:
    base = (base or "").rstrip("/")
    path = (path or "").lstrip("/")
    return f"{base}/{path}" if path else base


# -----------------------------
# DB routing (Admin-driven)
# -----------------------------
def pick_db_route(bot_name: str, canon: Dict) -> Optional[Dict]:
    """
    Look up CommandRoute for this bot + trigger/(optional) command.
    Returns: {"url": str, "headers": dict, "timeout": int}
             or None if no active route.
    """
    try:
        # Lazy import to avoid circulars at startup
        from .models import CommandRoute, Connection, TelegramBot  # noqa

        bot = TelegramBot.objects.filter(name=bot_name, is_active=True).first()
        if not bot:
            logger.info("router.no_bot bot=%s", bot_name)
            return None

        trigger = ((canon.get("command") or {}).get("trigger")) or "message"
        cmd = ((canon.get("command") or {}).get("name")) or None
        cmd = (cmd or "").strip().lstrip("/").lower() or None

        logger.info("router.lookup bot=%s trigger=%s cmd=%s", bot.username, trigger, cmd or "")

        qs = (
            CommandRoute.objects
            .select_related("connection")
            .filter(bot=bot, enabled=True, connection__is_active=True, trigger=trigger)
            .order_by("priority", "id")
        )
        # Log a small snapshot of candidates
        snapshot = list(qs.values("command_name", "trigger", "priority", "path", "connection__name")[:10])
        logger.info("router.candidates n=%s sample=%s", qs.count(), snapshot)

        # Prefer exact command; then default (blank/null)
        route = (
            qs.filter(command_name=cmd).first()
            or qs.filter(command_name__isnull=True).first()
            or qs.filter(command_name="").first()
        )
        if not route:
            logger.info("router.no_match bot=%s trigger=%s cmd=%s", bot.username, trigger, cmd or "")
            return None

        conn: Connection = route.connection
        url = _compose_url(conn.base_url, route.path or conn.path_default)

        ok, why = _is_allowed(url, conn.allowed_host_set())
        if not ok:
            logger.warning(
                "router.db.reject url=%s reason=%s allowed=%s route_conn=%s",
                url, why, conn.allowed_host_set(), conn.name
            )
            return None

        headers = {}
        headers.update(conn.auth_headers())
        headers.update(conn.extra_headers())

        logger.info(
            "router.route_ok bot=%s trigger=%s cmd=%s url=%s conn=%s timeout=%ss",
            bot.username, trigger, cmd or "", url, conn.name, conn.timeout_s
        )
        return {"url": url, "headers": headers, "timeout": conn.timeout_s}

    except Exception as e:
        logger.exception("router.db.error: %s", e)
        return None


# -----------------------------
# HTTP POST (DB routes)
# -----------------------------
def post_json(url: str, payload: Dict, timeout: float = 4.0, headers: Optional[Dict] = None) -> Tuple[int, Dict]:
    """
    Blocking POST JSON; never raises.
    Returns (status_code, body_json_or_wrapper).
    `headers` is optional for DB routes.
    """
    hdrs = {"Content-Type": "application/json", **(headers or {})}
    data = json.dumps(payload, ensure_ascii=False).encode("utf-8")
    t0 = time.perf_counter()

    if _HAS_REQUESTS:
        try:
            r = requests.post(url, data=data, headers=hdrs, timeout=timeout)
            dt = (time.perf_counter() - t0) * 1000
            try:
                body = r.json()
                preview = str(body)[:200].replace("\n", " ")
            except Exception:
                body = {"text": r.text[:2000]}
                preview = body["text"][:200].replace("\n", " ")
            logger.info("router.http requests url=%s status=%s t_ms=%.1f body~=%s", url, r.status_code, dt, preview)
            return r.status_code, body
        except Exception as e:
            dt = (time.perf_counter() - t0) * 1000
            logger.exception("router.requests_failed url=%s t_ms=%.1f err=%s", url, dt, e.__class__.__name__)
            return 502, {"ok": False, "error": f"requests_failed:{e.__class__.__name__}"}
    else:
        try:
            req = urllib.request.Request(url, data=data, headers=hdrs, method="POST")
            with urllib.request.urlopen(req, timeout=timeout) as resp:  # nosec
                raw = resp.read(65536)
                dt = (time.perf_counter() - t0) * 1000
                try:
                    body = json.loads(raw.decode("utf-8"))
                    preview = str(body)[:200].replace("\n", " ")
                except Exception:
                    txt = raw.decode("utf-8", errors="replace")[:2000]
                    body = {"text": txt}
                    preview = txt[:200].replace("\n", " ")
                status = getattr(resp, "status", 200)
                logger.info("router.http urllib url=%s status=%s t_ms=%.1f body~=%s", url, status, dt, preview)
                return status, body
        except Exception as e:
            dt = (time.perf_counter() - t0) * 1000
            logger.exception("router.urllib_failed url=%s t_ms=%.1f err=%s", url, dt, e.__class__.__name__)
            return 502, {"ok": False, "error": f"urllib_failed:{e.__class__.__name__}"}