124 lines
4.4 KiB
Python
124 lines
4.4 KiB
Python
# pxy_bots/router.py
|
|
import json
|
|
import logging
|
|
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:
|
|
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
|
|
|
|
qs = (
|
|
CommandRoute.objects
|
|
.select_related("connection")
|
|
.filter(bot=bot, enabled=True, connection__is_active=True, trigger=trigger)
|
|
.order_by("priority", "id")
|
|
)
|
|
|
|
# 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:
|
|
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", url, why, conn.allowed_host_set())
|
|
return None
|
|
|
|
headers = {}
|
|
headers.update(conn.auth_headers())
|
|
headers.update(conn.extra_headers())
|
|
|
|
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")
|
|
|
|
if _HAS_REQUESTS:
|
|
try:
|
|
r = requests.post(url, data=data, headers=hdrs, timeout=timeout)
|
|
try:
|
|
body = r.json()
|
|
except Exception:
|
|
body = {"text": r.text[:2000]}
|
|
return r.status_code, body
|
|
except Exception as e:
|
|
logger.exception("router.requests_failed url=%s", url)
|
|
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)
|
|
try:
|
|
body = json.loads(raw.decode("utf-8"))
|
|
except Exception:
|
|
body = {"text": raw.decode("utf-8", errors="replace")[:2000]}
|
|
return getattr(resp, "status", 200), body
|
|
except Exception as e:
|
|
logger.exception("router.urllib_failed url=%s", url)
|
|
return 502, {"ok": False, "error": f"urllib_failed:{e.__class__.__name__}"}
|