source_client.py

import json
import time
import random
import string
import uuid
import asyncio
import httpx
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, List, Optional
from urllib.parse import quote

from config import Config

JST = timezone(timedelta(hours=9))


class SourceClient:
    def __init__(self):
        self._cid = None
        self._cid_ts = 0
        self._auth = None
        self._auth_ts = 0
        self._channels = None
        self._channels_ts = 0
        self._auth_lock = None
        self._auth_lock_loop = None
        self._acct_idx = 0
        self._device_id = None
        self._pool = []
        self._pool_lock = None
        self._pool_lock_loop = None

    def _timeout(self, read=30.0):
        return httpx.Timeout(connect=10.0, read=read, write=10.0, pool=10.0)

    async def get_cid(self, force=False) -> str:
        if not force and self._cid and (time.time() - self._cid_ts) < 86400:
            return self._cid
        async with httpx.AsyncClient(timeout=self._timeout()) as client:
            r = await client.get(Config.get_cid_url())
            r.raise_for_status()
            data = r.json()
            if "cid" not in data:
                raise ValueError("cid missing in response")
            self._cid = data["cid"]
            self._cid_ts = time.time()
            return self._cid

    def _alock(self):
        loop = asyncio.get_running_loop()
        if self._auth_lock is None or self._auth_lock_loop is not loop:
            self._auth_lock = asyncio.Lock()
            self._auth_lock_loop = loop
        return self._auth_lock

    async def get_auth(self, force=False) -> Dict[str, Any]:
        if not force and self._auth and (time.time() - self._auth_ts) < 10800:
            return self._auth
        async with self._alock():
            ttl = 30 if force else 10800
            if self._auth and (time.time() - self._auth_ts) < ttl:
                return self._auth
            return await self._do_login()

    async def rotate_account(self) -> Dict[str, Any]:
        async with self._alock():
            if self._auth and (time.time() - self._auth_ts) < 30:
                return self._auth
            self._device_id = _gen_device_id(Config.LOGIN_DEVICE_ID)
            return await self._do_login()

    def _plock(self):
        loop = asyncio.get_running_loop()
        if self._pool_lock is None or self._pool_lock_loop is not loop:
            self._pool_lock = asyncio.Lock()
            self._pool_lock_loop = loop
        return self._pool_lock

    async def _login_device(self, device_id) -> Dict[str, Any]:
        cid = await self.get_cid()
        acct = {"password": Config.LOGIN_PASSWORD, "app_id": Config.LOGIN_APP_ID, "device_id": device_id}
        async with httpx.AsyncClient(timeout=self._timeout()) as client:
            r = await client.get(Config.get_login_url_for(cid, acct))
            r.raise_for_status()
            data = r.json()
            if data.get("code") != "OK":
                raise ValueError("login failed: " + str(data.get("message", "unknown")))
            product = json.loads(data.get("product_config", "{}"))
            tok = {
                "access_token": data["access_token"],
                "refresh_token": data.get("refresh_token", ""),
                "acct_cid": data.get("cid", ""),
                "app_id": Config.LOGIN_APP_ID,
                "device_id": device_id,
                "vms_host": str(product.get("vms_host", "")).rstrip("/"),
                "vms_uid": str(product.get("vms_uid", "")),
                "ts": time.time(),
                "rts": time.time(),
            }
            if not tok["access_token"]:
                raise ValueError("no access_token")
            return tok

    async def _refresh_one(self, tok) -> bool:
        if not tok.get("refresh_token"):
            return False
        url = Config.get_refresh_url(tok.get("refresh_token", ""), tok.get("acct_cid", ""),
                                     tok.get("app_id", ""), tok.get("device_id", ""))
        if not url:
            return False
        headers = {"Referer": Config.REQUIRED_REFERER, "User-Agent": "Mozilla/5.0"}
        async with httpx.AsyncClient(timeout=self._timeout()) as client:
            r = await client.get(url, headers=headers)
            r.raise_for_status()
            data = r.json()
            if data.get("code") != "OK" or not data.get("access_token"):
                return False
            tok["access_token"] = data["access_token"]
            if data.get("refresh_token"):
                tok["refresh_token"] = data["refresh_token"]
            tok["ts"] = time.time()
            return True

    async def build_pool(self, n) -> int:
        n = max(1, int(n))
        pool: List[Dict[str, Any]] = []
        for _ in range(n):
            try:
                pool.append(await self._login_device(_gen_device_id(Config.LOGIN_DEVICE_ID)))
            except Exception:
                pass
        if not pool:
            try:
                pool.append(await self._login_device(Config.LOGIN_DEVICE_ID or _gen_device_id("")))
            except Exception:
                pass
        self._pool = pool
        return len(pool)

    def pool_size(self) -> int:
        return len(self._pool)

    def token_at(self, i) -> str:
        p = self._pool
        if not p:
            return ""
        return p[i % len(p)].get("access_token", "")

    async def _fallback_token(self) -> str:
        try:
            auth = await self.get_auth()
            return auth.get("access_token", "")
        except Exception:
            return ""

    async def refresh_token_at(self, i) -> str:
        async with self._plock():
            p = self._pool
            if not p:
                return await self._fallback_token()
            tok = p[i % len(p)]
            if time.time() - tok.get("ts", 0) < 8:
                return tok.get("access_token", "")
            try:
                await self._refresh_one(tok)
            except Exception:
                pass
            return tok.get("access_token", "")

    async def rotate_pool_at(self, i) -> str:
        async with self._plock():
            p = self._pool
            if not p:
                return await self._fallback_token()
            if len(p) <= 1:
                tok = p[0]
                if time.time() - tok.get("rts", 0) >= 30:
                    tok["rts"] = time.time()
                    try:
                        await self._refresh_one(tok)
                    except Exception:
                        pass
                return tok.get("access_token", "")
            self._rr = (getattr(self, "_rr", i) + 1) % len(p)
            if self._rr == i % len(p):
                self._rr = (self._rr + 1) % len(p)
            return p[self._rr].get("access_token", "")

    async def _do_login(self, _retry=0) -> Dict[str, Any]:
        cid = await self.get_cid(force=(_retry > 0))
        if self._device_id:
            acct = {"password": Config.LOGIN_PASSWORD, "app_id": Config.LOGIN_APP_ID, "device_id": self._device_id}
            login_url = Config.get_login_url_for(cid, acct)
        else:
            accts = Config.login_accounts()
            if accts:
                acct = accts[self._acct_idx % len(accts)]
                login_url = Config.get_login_url_for(cid, acct)
            else:
                login_url = Config.get_login_url(cid)
        async with httpx.AsyncClient(timeout=self._timeout()) as client:
            r = await client.get(login_url)
            r.raise_for_status()
            data = r.json()
            if data.get("code") != "OK":
                msg = str(data.get("message", "unknown"))
                if "cid" in msg.lower() and _retry < 2:
                    return await self._do_login(_retry=_retry + 1)
                raise ValueError("login failed: " + msg)
            product = json.loads(data.get("product_config", "{}"))
            auth = {
                "access_token": data["access_token"],
                "vms_host": str(product["vms_host"]).rstrip("/"),
                "vms_uid": str(product["vms_uid"]),
            }
            if not all(auth.values()):
                raise ValueError("incomplete auth")
            self._auth = auth
            self._auth_ts = time.time()
            return auth

    async def get_channels(self, force=False) -> List[dict]:
        if not force and self._channels and (time.time() - self._channels_ts) < 86400:
            return self._channels
        auth = await self.get_auth()
        url = Config.get_list_url(auth["vms_uid"], with_epg=False)
        headers = {"Referer": Config.REQUIRED_REFERER, "User-Agent": "Mozilla/5.0"}
        async with httpx.AsyncClient(timeout=self._timeout()) as client:
            r = await client.get(url, headers=headers)
            if r.status_code in (401, 403):
                await self.get_auth(force=True)
                return await self.get_channels(force=True)
            r.raise_for_status()
            data = r.json()
            channels = [c for c in data.get("result", [])
                        if c.get("id") and c.get("no") and c.get("name") and c.get("playpath")]
            if not channels:
                raise ValueError("no channels")
            self._channels = channels
            self._channels_ts = time.time()
            return channels

    async def get_all_epg(self) -> Dict[str, List[dict]]:
        auth = await self.get_auth()
        url = Config.get_list_url(auth["vms_uid"], with_epg=True)
        headers = {"Referer": Config.REQUIRED_REFERER, "User-Agent": "Mozilla/5.0"}
        timeout = httpx.Timeout(connect=10.0, read=float(Config.EPG_READ_TIMEOUT), write=30.0, pool=10.0)
        async with httpx.AsyncClient(timeout=timeout) as client:
            r = await client.get(url, headers=headers)
            r.raise_for_status()
            data = r.json()
        result: Dict[str, List[dict]] = {}
        for ch in data.get("result", []):
            cid = ch.get("id")
            if not cid:
                continue
            raw = ch.get("record_epg")
            if not raw:
                result[str(cid)] = []
                continue
            try:
                epg = json.loads(raw)
            except Exception:
                result[str(cid)] = []
                continue
            programs = []
            for i, p in enumerate(epg):
                if not p.get("time"):
                    continue
                if not p.get("time_end"):
                    if i + 1 < len(epg) and epg[i + 1].get("time"):
                        p["time_end"] = epg[i + 1]["time"]
                    else:
                        continue
                programs.append(p)
            result[str(cid)] = programs
        return result

    def make_vod_m3u8(self, program: dict, auth: dict) -> str:
        path = program.get("path")
        if not path:
            return program.get("m3u8") or ""
        path = str(path).strip().lstrip("/")
        if not path.startswith("query/"):
            path = "query/" + path
        if path.endswith(".m3u8"):
            path = path[:-5]
        token = quote(auth.get("access_token", ""), safe="")
        host = Config.UPSTREAM_HOST_VOD.rstrip("/")
        if not host or not token:
            return ""
        return host + "/" + path + ".m3u8?type=vod&__cross_domain_user=" + token


def _gen_device_id(template) -> str:
    t = str(template or "").strip()
    if t and "-" in t and len(t.replace("-", "")) == 32:
        u = str(uuid.uuid4())
        return u.upper() if t == t.upper() else u
    digits = "0123456789abcdef"
    if t and all(c in digits for c in t.lower()):
        out = "".join(random.choice(digits) for _ in range(len(t)))
        return out.upper() if t == t.upper() else out
    if t and t.isalnum():
        pool = string.ascii_lowercase + string.digits
        return "".join(random.choice(pool) for _ in range(len(t)))
    return uuid.uuid4().hex


def normalize_ts(t) -> Optional[float]:
    if t is None:
        return None
    try:
        s = str(t).strip()
        if len(s) == 14 and s.isdigit():
            return datetime.strptime(s, "%Y%m%d%H%M%S").replace(tzinfo=JST).timestamp()
        v = float(s)
        return v / 1000.0 if v > 9999999999 else v
    except Exception:
        return None


source = SourceClient()