Resolve Kwai Cdn Url

[project]
name = "resolve-kwai-cdn-url"
version = "0.1.0"
description = "Resolve Kuaishou share links to CDN URLs using videodl (videofetch) or GraphQL fallback."
requires-python = ">=3.9"
dependencies = [
  "videofetch @ git+https://github.com/debugtalk/videodl@master",
  "tls-client",
]

#!/usr/bin/env python3
"""Shared helpers for Kwai/Kuaishou CSV+JSONL tooling."""

from __future__ import annotations

import json
import os
import tempfile
from typing import Dict, Optional, Set, Tuple

def classify_bad_cdn_url(url: str) -> Optional[str]:
    if not url:
        return None
    lower = url.strip().lower()
    if lower.startswith("https://live.kuaishou.com") or lower.startswith(
        "http://live.kuaishou.com"
    ):
        return "live.kuaishou.com is not a CDN url"
    if lower.startswith("https://captcha.zt.kuaishou.com") or lower.startswith(
        "http://captcha.zt.kuaishou.com"
    ):
        return f"captcha url is not a CDN url: {url}"
    return None


def is_success_record(
    item: Dict[str, object],
    cdn_field: str = "CDNURL",
    error_field: Optional[str] = "error_msg",
) -> bool:
    cdn = item.get(cdn_field) or item.get("CDNURL") or item.get("cdn_url") or item.get("cdnUrl")
    err = ""
    if error_field:
        err = item.get(error_field, "")
    elif "error_msg" in item:
        err = item.get("error_msg", "")
    if err is None:
        err = ""
    if isinstance(err, str):
        err = err.strip()
    return bool(isinstance(cdn, str) and cdn.strip()) and not err


def extract_url_from_record(item: Dict[str, object], url_field: str) -> str:
    for key in (url_field, "URL", "url"):
        val = item.get(key)
        if isinstance(val, str) and val.strip():
            return val.strip()
    return ""


def load_resume_success_urls(
    output_path: str,
    url_field: str,
    cdn_field: str = "CDNURL",
    error_field: Optional[str] = "error_msg",
) -> Set[str]:
    resolved: Set[str] = set()
    try:
        with open(output_path, "r", encoding="utf-8") as handle:
            for line in handle:
                line = line.strip()
                if not line:
                    continue
                try:
                    item = json.loads(line)
                except json.JSONDecodeError:
                    continue
                if not isinstance(item, dict):
                    continue
                if not is_success_record(item, cdn_field=cdn_field, error_field=error_field):
                    continue
                url = extract_url_from_record(item, url_field)
                if url:
                    resolved.add(url)
    except FileNotFoundError:
        return set()
    return resolved


def clean_output_jsonl(
    output_path: str,
    cdn_field: str = "CDNURL",
    error_field: Optional[str] = "error_msg",
) -> Tuple[int, int]:
    kept = 0
    removed = 0
    try:
        with open(output_path, "r", encoding="utf-8") as handle:
            dir_name = os.path.dirname(output_path) or "."
            with tempfile.NamedTemporaryFile(
                mode="w", encoding="utf-8", delete=False, dir=dir_name
            ) as tmp:
                for line in handle:
                    raw = line.strip()
                    if not raw:
                        continue
                    try:
                        item = json.loads(raw)
                    except json.JSONDecodeError:
                        removed += 1
                        continue
                    if not isinstance(item, dict):
                        removed += 1
                        continue
                    if is_success_record(item, cdn_field=cdn_field, error_field=error_field):
                        tmp.write(json.dumps(item, ensure_ascii=False) + "\n")
                        kept += 1
                    else:
                        removed += 1
        os.replace(tmp.name, output_path)
    except FileNotFoundError:
        return (0, 0)
    return (kept, removed)


def load_cookie_value(
    cookie: Optional[str] = None,
    cookie_file: Optional[str] = None,
) -> Optional[str]:
    if cookie_file:
        try:
            cookie_text = open(cookie_file, "r", encoding="utf-8").read().strip()
        except Exception:
            cookie_text = ""
        if cookie_text:
            try:
                data = json.loads(cookie_text)
            except Exception:
                data = None
            if isinstance(data, dict) and "cookies" in data:
                data = data["cookies"]
            if isinstance(data, list):
                pairs = []
                for item in data:
                    if not isinstance(item, dict):
                        continue
                    name = item.get("name")
                    value = item.get("value")
                    if isinstance(name, str) and isinstance(value, str) and name:
                        pairs.append(f"{name}={value}")
                return "; ".join(pairs) if pairs else None
            return cookie_text
        return None
    return cookie

#!/usr/bin/env python3
"""Read CSV, extract Kuaishou CDN URLs, and write JSONL."""

import argparse
import csv
import hashlib
import json
import os
import signal
import sys
import time
from collections import Counter
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Dict, List, Set, Tuple

from kwai_extract_cdn import extract_cdn_url_detail
from kwai_common import classify_bad_cdn_url, clean_output_jsonl, load_resume_success_urls

CLIENT_IDENTIFIERS = (
    "chrome_120",
    "firefox_120",
    "safari_16_0",
)


def _write_jsonl_row(row: Dict[str, str], output) -> None:
    output.write(json.dumps(row, ensure_ascii=False))
    output.write("\n")


def _process_one(
    share_url: str,
    cookie: str,
    cookie_file: str,
    proxy: str,
    endpoints: List[str],
    timeout: int,
    jitter: Tuple[float, float],
) -> Tuple[str, str]:
    if not share_url:
        return "", "empty url"
    client_identifier = _select_client_identifier(share_url)
    return extract_cdn_url_detail(
        share_url,
        client_identifier=client_identifier,
        cookie=cookie,
        cookie_file=cookie_file,
        proxy=proxy,
        endpoints=endpoints or None,
        timeout=timeout,
        jitter=jitter,
    )


def _select_client_identifier(share_url: str) -> str:
    if not CLIENT_IDENTIFIERS:
        return "chrome_120"
    digest = hashlib.md5(share_url.encode("utf-8")).digest()
    idx = int.from_bytes(digest[:2], "big") % len(CLIENT_IDENTIFIERS)
    return CLIENT_IDENTIFIERS[idx]


def _normalize_cdn_url(cdn_url: str, err: str) -> Tuple[str, str]:
    if not cdn_url:
        return "", err
    bad_reason = classify_bad_cdn_url(cdn_url)
    if bad_reason:
        if err and err != bad_reason:
            return "", f"{err}; {bad_reason}"
        return "", bad_reason
    return cdn_url, err


def _bucket_failure(err: str) -> str:
    if not err:
        return "unknown"
    lower = err.lower()
    if "captcha url is not a cdn url" in lower:
        return "captcha_url"
    if "live.kuaishou.com is not a cdn url" in lower:
        return "live_url"
    if "photoid not found" in lower:
        return "photo_id_missing"
    if "cdn url not found" in lower:
        return "cdn_not_found"
    if "empty url" in lower:
        return "empty_url"
    return "other"


def _emit_stats(
    processed: int,
    success_total: int,
    failure_total: int,
    skipped: int,
    failure_buckets: Counter[str],
) -> None:
    if processed <= 0:
        return
    bucket_parts = ", ".join(
        f"{key}={value}" for key, value in failure_buckets.items()
    )
    sys.stderr.write(
        "stats: processed="
        f"{processed} success={success_total} failed={failure_total}"
        f" skipped={skipped} buckets={{ {bucket_parts} }}\n"
    )


def main() -> int:
    parser = argparse.ArgumentParser(description="CSV -> JSONL with Kuaishou CDN URLs")
    parser.add_argument("csv", help="Input CSV path")
    parser.add_argument("--url-col", default="URL", help="Column name that holds the Kuaishou URL")
    parser.add_argument("--cdn-col", default="CDNURL", help="Output column name for CDN URL")
    parser.add_argument(
        "--error-col",
        default="error_msg",
        help="Output column for error message (default: error_msg). Use empty string to disable.",
    )
    parser.add_argument("--cookie", default=None, help="Raw Cookie header value")
    parser.add_argument(
        "--cookie-file",
        default=None,
        help="Path to a cookie.txt file or JSON cookie array",
    )
    parser.add_argument(
        "--proxy",
        default=None,
        help="Proxy URL, e.g. http://user:pass@host:port",
    )
    parser.add_argument(
        "--endpoint",
        action="append",
        default=None,
        help="GraphQL endpoint override (can be provided multiple times)",
    )
    parser.add_argument("--sleep", type=float, default=0.0, help="Sleep between requests (seconds)")
    parser.add_argument("--workers", type=int, default=1, help="Concurrent workers")
    parser.add_argument("--progress-every", type=int, default=10, help="Progress log interval")
    parser.add_argument(
        "--jitter",
        default="0,0",
        help="Sleep per request as min,max seconds (e.g., 1,3)",
    )
    parser.add_argument("--timeout", type=int, default=10, help="Request timeout (seconds)")
    parser.add_argument("--output", default="-", help="Output JSONL path or '-' for stdout")
    parser.add_argument(
        "--resume",
        action="store_true",
        help="Skip URLs already resolved in output JSONL (CDNURL set, error empty)",
    )
    args = parser.parse_args()

    output_mode = "w"
    resume_urls: Set[str] = set()
    if args.resume:
        if args.output == "-":
            sys.stderr.write("--resume requires --output file\n")
            return 2
        kept, removed = clean_output_jsonl(
            args.output,
            cdn_field=args.cdn_col,
            error_field=args.error_col or None,
        )
        if removed:
            sys.stderr.write(f"resume: removed {removed} failed rows\n")
        resume_urls = load_resume_success_urls(
            args.output,
            args.url_col,
            cdn_field=args.cdn_col,
            error_field=args.error_col or None,
        )
        output_mode = "a"

    output = sys.stdout if args.output == "-" else open(args.output, output_mode, encoding="utf-8")
    jitter_vals = (0.0, 0.0)
    try:
        parts = [p.strip() for p in args.jitter.split(",")]
        if len(parts) == 2:
            jitter_vals = (float(parts[0]), float(parts[1]))
    except Exception:
        jitter_vals = (0.0, 0.0)

    with open(args.csv, "r", encoding="utf-8-sig", newline="") as f:
        reader = csv.DictReader(f)
        if not reader.fieldnames or args.url_col not in reader.fieldnames:
            sys.stderr.write(
                f"Missing URL column '{args.url_col}'. Available: {reader.fieldnames}\n"
            )
            return 2

        rows = list(reader)
        total = len(rows)
        skipped = 0
        if resume_urls:
            rows = [
                row
                for row in rows
                if row.get(args.url_col, "").strip() not in resume_urls
            ]
            skipped = total - len(rows)
            if skipped:
                sys.stderr.write(f"resume: skipped {skipped} already resolved\n")
        if not rows:
            sys.stderr.write("resume: all rows already resolved\n")
            return 0
        total_remaining = len(rows)
        completed = skipped
        processed = 0
        success_total = 0
        failure_total = 0
        success_batch = 0
        failure_buckets: Counter[str] = Counter()

        try:
            if args.workers <= 1:
                for row in rows:
                    share_url = row.get(args.url_col, "")
                    cdn_url, err = _process_one(
                        share_url,
                        cookie=args.cookie,
                        cookie_file=args.cookie_file,
                        proxy=args.proxy,
                        endpoints=args.endpoint or [],
                        timeout=args.timeout,
                        jitter=jitter_vals,
                    )
                    cdn_url, err = _normalize_cdn_url(cdn_url, err)
                    row[args.cdn_col] = cdn_url or ""
                    if args.error_col:
                        row[args.error_col] = err
                    _write_jsonl_row(row, output)
                    completed += 1
                    processed += 1
                    if cdn_url:
                        success_batch += 1
                        success_total += 1
                    else:
                        failure_total += 1
                        failure_buckets[_bucket_failure(err)] += 1
                    if args.progress_every and completed % args.progress_every == 0:
                        sys.stderr.write(
                            f"progress: {completed}/{total} success: {success_batch}\n"
                        )
                        success_batch = 0
                    if args.sleep:
                        time.sleep(args.sleep)
            else:
                executor = ThreadPoolExecutor(max_workers=args.workers)
                try:
                    future_map = {}
                    for idx, row in enumerate(rows):
                        share_url = row.get(args.url_col, "")
                        future = executor.submit(
                            _process_one,
                            share_url,
                            args.cookie,
                            args.cookie_file,
                            args.proxy,
                            args.endpoint or [],
                            args.timeout,
                            jitter_vals,
                        )
                        future_map[future] = idx

                    results: List[Tuple[str, str]] = [("", "")] * total_remaining
                    ready: Dict[int, Tuple[str, str]] = {}
                    next_idx = 0
                    for future in as_completed(future_map):
                        idx = future_map[future]
                        try:
                            results[idx] = _normalize_cdn_url(*future.result())
                        except Exception as exc:
                            results[idx] = ("", str(exc))
                        ready[idx] = results[idx]
                        completed += 1
                        if results[idx][0]:
                            success_batch += 1
                        if args.progress_every and completed % args.progress_every == 0:
                            sys.stderr.write(
                                f"progress: {completed}/{total} success: {success_batch}\n"
                            )
                            success_batch = 0

                        while next_idx in ready:
                            cdn_url, err = ready.pop(next_idx)
                            row = rows[next_idx]
                            row[args.cdn_col] = cdn_url or ""
                            if args.error_col:
                                row[args.error_col] = err
                            _write_jsonl_row(row, output)
                            processed += 1
                            if cdn_url:
                                success_total += 1
                            else:
                                failure_total += 1
                                failure_buckets[_bucket_failure(err)] += 1
                            next_idx += 1
                finally:
                    executor.shutdown(wait=False, cancel_futures=True)
        except KeyboardInterrupt:
            sys.stderr.write("Interrupted by user, exiting gracefully.\n")
            if args.workers > 1:
                signal.signal(signal.SIGINT, signal.SIG_IGN)
            _emit_stats(processed, success_total, failure_total, skipped, failure_buckets)
            sys.stderr.flush()
            os._exit(130)
        finally:
            _emit_stats(processed, success_total, failure_total, skipped, failure_buckets)

    if output is not sys.stdout:
        output.close()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

#!/usr/bin/env python3
"""Extract Kuaishou CDN URL from a share link."""

import argparse
import json
import random
import re
import sys
import time
from typing import Any, Dict, Iterable, List, Optional, Tuple
from urllib.parse import parse_qs, urlparse

import tls_client

from kwai_common import classify_bad_cdn_url, load_cookie_value

DEFAULT_UA = (
    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) "
    "AppleWebKit/537.36 (KHTML, like Gecko) "
    "Chrome/120.0.0.0 Safari/537.36"
)

MOBILE_UA = (
    "Mozilla/5.0 (iPhone; CPU iPhone OS 16_6 like Mac OS X) "
    "AppleWebKit/605.1.15 (KHTML, like Gecko) Version/16.6 "
    "Mobile/15E148 Safari/604.1"
)

DEFAULT_ENDPOINTS = [
    "https://www.kuaishou.com/graphql",
    "https://live.kuaishou.com/m_graphql",
]

VISION_VIDEO_DETAIL_QUERY = (
    "query visionVideoDetail($photoId: String) {"
    "  visionVideoDetail(photoId: $photoId) {"
    "    photo {"
    "      photoUrl"
    "      mainNoWatermarkUrl"
    "      mainUrl"
    "      videoResource {"
    "        h265 { adaptSetRepresentation }"
    "        h264 { adaptSetRepresentation }"
    "      }"
    "      manifest"
    "    }"
    "  }"
    "}"
)

URL_FIELD_HINTS = (
    "photoUrl",
    "mainNoWatermarkUrl",
    "mainUrl",
    "playUrl",
    "srcNoMark",
    "videoUrl",
)

DOMAIN_HINTS = (
    "kuaishou",
    "kwaicdn",
    "ks-cdn",
    "ksyuncdn",
    "gifshow",
)


PHOTO_ID_RE = re.compile(r"/(?:photo|short-video)/([^/?#]+)")


def extract_first_url(text: str) -> Optional[str]:
    match = re.search(r"https?://[^\s\"'<>]+", text)
    if not match:
        return None
    return match.group(0).rstrip(").,;]}")


def _collect_urls(obj: Any, out: List[str]) -> None:
    if isinstance(obj, dict):
        for k, v in obj.items():
            if isinstance(v, str):
                if v.startswith("http"):
                    out.append(v)
            else:
                _collect_urls(v, out)
    elif isinstance(obj, list):
        for item in obj:
            _collect_urls(item, out)


def _rank_url(url: str) -> Tuple[int, int]:
    score = 0
    if any(h in url for h in DOMAIN_HINTS):
        score += 10
    if url.endswith(".mp4"):
        score += 5
    return (score, len(url))


def _pick_best_url(urls: Iterable[str]) -> Optional[str]:
    unique = list(
        dict.fromkeys(
            u for u in urls if u.startswith("http") and not classify_bad_cdn_url(u)
        )
    )
    if not unique:
        return None
    return sorted(unique, key=_rank_url, reverse=True)[0]


def _extract_photo_id(url: str) -> Optional[str]:
    match = PHOTO_ID_RE.search(url)
    if match:
        return match.group(1)
    parsed = urlparse(url)
    if parsed.path:
        parts = [p for p in parsed.path.split("/") if p]
        if parts:
            if parts[-1] and len(parts[-1]) >= 6:
                if parts[-2:] and parts[-2] in {"short-video", "photo", "video"}:
                    return parts[-1]
    qs = parse_qs(parsed.query)
    for key in ("photoId", "shareObjectId", "videoId", "shortVideoId"):
        if key in qs and qs[key]:
            return qs[key][0]
    return None


def _resolve_url(
    session: tls_client.Session,
    url: str,
    timeout: int = 10,
    proxy: Optional[str] = None,
) -> str:
    try:
        resp = session.get(
            url,
            allow_redirects=True,
            timeout_seconds=timeout,
            proxy=proxy,
        )
        return resp.url
    except Exception:
        return url


def _parse_init_state(html: str) -> Optional[Dict[str, Any]]:
    marker = "window.INIT_STATE="
    idx = html.find(marker)
    if idx < 0:
        return None
    start = idx + len(marker)
    depth = 0
    end = None
    for i in range(start, len(html)):
        ch = html[i]
        if ch == "{":
            depth += 1
        elif ch == "}":
            depth -= 1
            if depth == 0:
                end = i + 1
                break
    if end is None:
        return None
    payload = html[start:end]
    try:
        return json.loads(payload)
    except Exception:
        return None


def _search_dict_by_key(obj: Any, key: str) -> Iterable[Any]:
    if isinstance(obj, dict):
        for k, v in obj.items():
            if k == key:
                yield v
            if isinstance(v, dict) or isinstance(v, list):
                yield from _search_dict_by_key(v, key)
    elif isinstance(obj, list):
        for item in obj:
            yield from _search_dict_by_key(item, key)


def _extract_from_init_state(html: str) -> Optional[str]:
    state = _parse_init_state(html)
    if not state:
        return None
    reps = list(_search_dict_by_key(state, "representation"))
    if not reps:
        return None
    best = None
    best_score = -1
    for rep_list in reps:
        if not isinstance(rep_list, list):
            continue
        for rep in rep_list:
            if not isinstance(rep, dict):
                continue
            candidate = rep.get("url") or rep.get("cdnUrl") or rep.get("downloadUrl")
            if not candidate:
                continue
            size = rep.get("fileSize") or 0
            if size > best_score:
                best_score = size
                best = candidate
    return best


def _graphql_fetch(
    session: tls_client.Session,
    endpoint: str,
    photo_id: str,
    timeout: int,
    proxy: Optional[str] = None,
) -> Optional[Dict[str, Any]]:
    payload = {
        "operationName": "visionVideoDetail",
        "query": VISION_VIDEO_DETAIL_QUERY,
        "variables": {"photoId": photo_id},
    }
    try:
        resp = session.post(
            endpoint,
            json=payload,
            timeout_seconds=timeout,
            proxy=proxy,
        )
        if resp.status_code != 200:
            return None
        return resp.json()
    except Exception:
        return None


def _extract_from_graphql_payload(payload: Dict[str, Any]) -> Optional[str]:
    if not payload:
        return None
    data = payload.get("data") or {}
    detail = data.get("visionVideoDetail") or data.get("VisionVideoDetail") or {}
    photo = detail.get("photo") or detail.get("Photo") or {}
    urls: List[str] = []

    for field in URL_FIELD_HINTS:
        value = photo.get(field)
        if isinstance(value, str):
            urls.append(value)
        elif isinstance(value, list):
            urls.extend([v for v in value if isinstance(v, str)])

    if not urls:
        _collect_urls(payload, urls)

    return _pick_best_url(urls)


def _extract_from_html(html: str) -> Optional[str]:
    for pattern in (
        r"__APOLLO_STATE__\s*=\s*({.+?})\s*;",
        r"__NEXT_DATA__\s*=\s*({.+?})\s*</script>",
        r"__INITIAL_STATE__\s*=\s*({.+?})\s*;",
    ):
        match = re.search(pattern, html, re.DOTALL)
        if not match:
            continue
        try:
            data = json.loads(match.group(1))
        except Exception:
            continue
        urls: List[str] = []
        _collect_urls(data, urls)
        best = _pick_best_url(urls)
        if best:
            return best
    return None


def extract_cdn_url_detail(
    share_url: str,
    client_identifier: Optional[str] = None,
    cookie: Optional[str] = None,
    cookie_file: Optional[str] = None,
    proxy: Optional[str] = None,
    endpoints: Optional[List[str]] = None,
    timeout: int = 10,
    jitter: Tuple[float, float] = (0.0, 0.0),
) -> Tuple[Optional[str], str]:
    session_kwargs = {
        "client_identifier": client_identifier or "chrome_120",
        "random_tls_extension_order": True,
    }
    try:
        session = tls_client.Session(**session_kwargs)
    except Exception:
        session = tls_client.Session(
            client_identifier="chrome_120",
            random_tls_extension_order=True,
        )
    session.headers.update(
        {
            "User-Agent": DEFAULT_UA,
            "Referer": "https://www.kuaishou.com/",
            "Origin": "https://www.kuaishou.com",
            "Accept-Language": "zh-CN,zh;q=0.9,en;q=0.8",
        }
    )

    cookie = load_cookie_value(cookie=cookie, cookie_file=cookie_file)
    if cookie:
        session.headers["Cookie"] = cookie

    raw_url = extract_first_url(share_url) or share_url
    if jitter and (jitter[0] or jitter[1]):
        time.sleep(random.uniform(jitter[0], jitter[1]))

    resolved = _resolve_url(session, raw_url, timeout=timeout, proxy=proxy)
    photo_id = _extract_photo_id(resolved)
    if not photo_id:
        photo_id = _extract_photo_id(raw_url)
    if not photo_id:
        return None, "photoId not found"

    for endpoint in (endpoints or DEFAULT_ENDPOINTS):
        payload = _graphql_fetch(session, endpoint, photo_id, timeout, proxy=proxy)
        url = _extract_from_graphql_payload(payload or {})
        if url:
            bad_reason = classify_bad_cdn_url(url)
            if bad_reason:
                return None, bad_reason
            return url, ""

    try:
        mobile_url = f"https://m.kuaishou.com/fw/photo/{photo_id}"
        resp = session.get(
            mobile_url,
            headers={
                "User-Agent": MOBILE_UA,
                "Referer": "https://m.kuaishou.com/",
            },
            timeout_seconds=timeout,
            proxy=proxy,
            allow_redirects=True,
        )
        if resp.status_code == 200:
            url = _extract_from_init_state(resp.text)
            if url:
                bad_reason = classify_bad_cdn_url(url)
                if bad_reason:
                    return None, bad_reason
                return url, ""
    except Exception:
        pass

    try:
        resp = session.get(resolved, timeout_seconds=timeout, proxy=proxy)
        if resp.status_code == 200:
            url = _extract_from_html(resp.text)
            if url:
                bad_reason = classify_bad_cdn_url(url)
                if bad_reason:
                    return None, bad_reason
                return url, ""
    except Exception:
        pass

    return None, "cdn url not found"


def extract_cdn_url(
    share_url: str,
    client_identifier: Optional[str] = None,
    cookie: Optional[str] = None,
    cookie_file: Optional[str] = None,
    proxy: Optional[str] = None,
    endpoints: Optional[List[str]] = None,
    timeout: int = 10,
) -> Optional[str]:
    url, _ = extract_cdn_url_detail(
        share_url,
        client_identifier=client_identifier,
        cookie=cookie,
        cookie_file=cookie_file,
        proxy=proxy,
        endpoints=endpoints,
        timeout=timeout,
    )
    return url


def main() -> int:
    parser = argparse.ArgumentParser(description="Extract Kuaishou CDN URL from a share link")
    parser.add_argument("url", help="Kuaishou share URL")
    parser.add_argument("--cookie", default=None, help="Raw Cookie header value")
    parser.add_argument(
        "--cookie-file",
        default=None,
        help="Path to a cookie.txt file or JSON cookie array",
    )
    parser.add_argument(
        "--proxy",
        default=None,
        help="Proxy URL, e.g. http://user:pass@host:port",
    )
    parser.add_argument(
        "--endpoint",
        action="append",
        default=None,
        help="GraphQL endpoint override (can be provided multiple times)",
    )
    parser.add_argument(
        "--jitter",
        default="0,0",
        help="Sleep between requests as min,max seconds (e.g., 1,3)",
    )
    parser.add_argument("--timeout", type=int, default=10, help="Request timeout (seconds)")
    args = parser.parse_args()

    jitter_vals = (0.0, 0.0)
    try:
        parts = [p.strip() for p in args.jitter.split(",")]
        if len(parts) == 2:
            jitter_vals = (float(parts[0]), float(parts[1]))
    except Exception:
        jitter_vals = (0.0, 0.0)

    url, err = extract_cdn_url_detail(
        args.url,
        cookie=args.cookie,
        cookie_file=args.cookie_file,
        proxy=args.proxy,
        endpoints=args.endpoint,
        timeout=args.timeout,
        jitter=jitter_vals,
    )

    if not url:
        print("", end="")
        return 2

    print(url)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

#!/usr/bin/env python3
"""
Resolve Kuaishou share links to CDN URLs using videodl (videofetch).
Output JSONL: {"url": "...", "cdn_url": "...", "error_msg": "..."}
"""

from __future__ import annotations

import argparse
import csv
import json
import re
import signal
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Iterable, List, Optional, Sequence, Tuple
from urllib.parse import urlparse

import requests
from videodl import videodl as videodl_lib

from kwai_common import classify_bad_cdn_url, clean_output_jsonl, load_cookie_value, load_resume_success_urls

UA_MOBILE = (
    "Mozilla/5.0 (iPhone; CPU iPhone OS 16_6 like Mac OS X) "
    "AppleWebKit/605.1.15 (KHTML, like Gecko) Version/16.6 Mobile/15E148 Safari/604.1"
)
UA_DESKTOP = (
    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) "
    "AppleWebKit/537.36 (KHTML, like Gecko) "
    "Chrome/120.0.0.0 Safari/537.36"
)

_INTERRUPTED = False


def _handle_sigint(signum, frame) -> None:
    del signum, frame
    global _INTERRUPTED
    _INTERRUPTED = True
    signal.signal(signal.SIGINT, signal.SIG_IGN)


def extract_first_url(text: str) -> Optional[str]:
    match = re.search(r"https?://[^\s]+", text)
    if not match:
        return None
    url = match.group(0)
    return url.strip(" \t\r\n\"'<>[](){}，。；;:!?")

def build_proxies(proxy: Optional[str]) -> Optional[dict]:
    if not proxy:
        return None
    scheme = urlparse(proxy).scheme.lower()
    if scheme == "http":
        return {"http": proxy}
    if scheme == "https":
        return {"https": proxy}
    return {"http": proxy, "https": proxy}


def resolve_short_url(url: str, timeout: float, proxy: Optional[str]) -> str:
    if "v.kuaishou.com" not in url:
        return url
    try:
        proxies = build_proxies(proxy)
        resp = requests.get(
            url,
            allow_redirects=True,
            timeout=timeout,
            headers={"User-Agent": UA_MOBILE},
            proxies=proxies,
        )
        return resp.url or url
    except requests.RequestException:
        return url


def fetch_anonymous_cookie(
    timeout: float, proxy: Optional[str]
) -> Optional[str]:
    proxies = build_proxies(proxy)
    session = requests.Session()
    headers = {
        "User-Agent": UA_DESKTOP,
        "Referer": "https://www.kuaishou.com/",
        "Origin": "https://www.kuaishou.com",
        "Accept-Language": "zh-CN,zh;q=0.9,en;q=0.8",
    }
    try:
        session.get(
            "https://www.kuaishou.com/",
            timeout=timeout,
            headers=headers,
            proxies=proxies,
        )
    except requests.RequestException:
        return None
    if not session.cookies:
        return None
    pairs = [f"{c.name}={c.value}" for c in session.cookies]
    return "; ".join(pairs) if pairs else None


def detect_unavailable_reason(
    url: str, timeout: float, proxy: Optional[str]
) -> Optional[str]:
    try:
        proxies = build_proxies(proxy)
        resp = requests.get(
            url,
            allow_redirects=True,
            timeout=timeout,
            headers={"User-Agent": UA_MOBILE},
            proxies=proxies,
        )
    except requests.RequestException as exc:
        return f"http error: {exc}"
    if not resp.text:
        return None
    text = resp.text
    keywords = [
        ("找不到该作品", "not found: removed or unavailable"),
        ("作品已失效", "not found: removed or unavailable"),
        ("该作品已被删除", "deleted by author"),
        ("已下架", "removed from shelf"),
        ("内容不可见", "content not visible"),
        ("内容不存在", "content not found"),
    ]
    for needle, message in keywords:
        if needle in text:
            return message
    return None


def pick_download_url(info: dict) -> Optional[str]:
    for key in ("download_url", "download_urls", "url", "urls"):
        val = info.get(key)
        if isinstance(val, str) and val:
            return val
        if isinstance(val, list):
            for item in val:
                if isinstance(item, str) and item:
                    return item
                if isinstance(item, dict):
                    for subkey in ("download_url", "download_urls", "url", "urls"):
                        subval = item.get(subkey)
                        if isinstance(subval, str) and subval:
                            return subval
                        if isinstance(subval, list):
                            for subitem in subval:
                                if isinstance(subitem, str) and subitem:
                                    return subitem
    return None


def resolve_cdn_url(url: str, request_overrides: Optional[dict]) -> str:
    video_client = videodl_lib.VideoClient(
        allowed_video_sources=["KuaishouVideoClient"],
        requests_overrides={"KuaishouVideoClient": request_overrides or {}},
    )
    video_infos = video_client.parsefromurl(url)
    if not video_infos:
        return ""
    for info in video_infos:
        cdn = pick_download_url(info)
        if cdn:
            return cdn
    return ""


def process_one(
    text: str,
    timeout: float,
    proxy: Optional[str],
    cookie: Optional[str],
    cookie_file: Optional[str],
    anonymous_cookie: Optional[str],
) -> dict:
    raw = text.strip()
    result = {"url": "", "cdn_url": "", "error_msg": ""}
    if not raw:
        result["error_msg"] = "empty input"
        return result

    url = extract_first_url(raw) or raw
    result["url"] = url
    if not url.startswith("http"):
        result["error_msg"] = "no http/https url found"
        return result

    resolved = resolve_short_url(url, timeout=timeout, proxy=proxy)
    request_overrides = {"timeout": timeout}
    if proxy:
        request_overrides["proxies"] = build_proxies(proxy)
    headers = dict(request_overrides.get("headers") or {})
    headers.setdefault("User-Agent", UA_MOBILE)
    headers.setdefault("Referer", "https://www.kuaishou.com/")
    headers.setdefault("Origin", "https://www.kuaishou.com")
    headers.setdefault("Accept-Language", "zh-CN,zh;q=0.9,en;q=0.8")
    cookie_value = load_cookie_value(cookie=cookie, cookie_file=cookie_file)
    if not cookie_value and anonymous_cookie:
        cookie_value = anonymous_cookie
    if cookie_value:
        headers["Cookie"] = cookie_value
    if headers:
        request_overrides["headers"] = headers
    try:
        cdn_url = resolve_cdn_url(resolved, request_overrides)
        if not cdn_url:
            reason = detect_unavailable_reason(
                resolved, timeout=timeout, proxy=proxy
            )
            result["error_msg"] = reason or "cdn url not found in videodl response"
        else:
            bad_reason = classify_bad_cdn_url(cdn_url)
            if bad_reason:
                result["error_msg"] = bad_reason
            else:
                result["cdn_url"] = cdn_url
    except Exception as exc:  # pylint: disable=broad-except
        result["error_msg"] = str(exc)

    return result


def load_inputs(input_text: Optional[str]) -> List[str]:
    items: List[str] = []
    if input_text:
        items.append(input_text)
    return items


def write_jsonl(
    items: Iterable[dict], output_path: Optional[str], append: bool = False
) -> None:
    if output_path:
        mode = "a" if append else "w"
        out = open(output_path, mode, encoding="utf-8")
    else:
        out = sys.stdout
    try:
        for item in items:
            out.write(json.dumps(item, ensure_ascii=False) + "\n")
    finally:
        if output_path:
            out.close()


def load_csv_rows(
    csv_path: str,
    url_field: str,
) -> Tuple[List[dict], List[str], str]:
    with open(csv_path, "r", encoding="utf-8-sig", newline="") as handle:
        reader = csv.reader(handle)
        try:
            header = next(reader)
        except StopIteration:
            return [], [], url_field

        header = [col.strip() for col in header]
        kept_indices = [idx for idx, col in enumerate(header) if col]
        kept_headers = [header[idx] for idx in kept_indices]
        if not kept_headers:
            return [], [], url_field

        url_lookup = {name.lower(): name for name in kept_headers}
        if url_field not in kept_headers:
            mapped = url_lookup.get(url_field.lower())
            if mapped:
                url_field = mapped
            else:
                raise ValueError(f"URL field '{url_field}' not found in CSV header")

        rows: List[dict] = []
        for row in reader:
            if not row or not any(cell.strip() for cell in row):
                continue
            record: dict = {}
            for idx, name in zip(kept_indices, kept_headers):
                record[name] = row[idx].strip() if idx < len(row) else ""
            rows.append(record)

    return rows, kept_headers, url_field


def process_batch(
    batch: List[str],
    timeout: float,
    proxy: Optional[str],
    cookie: Optional[str],
    cookie_file: Optional[str],
    workers: int,
    progress_every: int,
    completed_offset: int,
    total: int,
    anonymous_cookie: Optional[str],
) -> List[dict]:
    global _INTERRUPTED
    if workers <= 1 or len(batch) == 1:
        results = []
        for idx, item in enumerate(batch, start=1):
            if _INTERRUPTED:
                break
            results.append(
                process_one(
                    item,
                    timeout=timeout,
                    proxy=proxy,
                    cookie=cookie,
                    cookie_file=cookie_file,
                    anonymous_cookie=anonymous_cookie,
                )
            )
            if progress_every > 0 and idx % progress_every == 0:
                done = completed_offset + idx
                print(f"progress: {done}/{total}", file=sys.stderr)
        if _INTERRUPTED:
            for item in batch[len(results) :]:
                results.append(
                    {"url": item, "cdn_url": "", "error_msg": "interrupted"}
                )
        return results

    results: List[Optional[dict]] = [None] * len(batch)
    completed = 0
    executor = ThreadPoolExecutor(max_workers=workers)
    try:
        future_map = {
            executor.submit(
                process_one,
                item,
                timeout,
                proxy,
                cookie,
                cookie_file,
                anonymous_cookie,
            ): idx
            for idx, item in enumerate(batch)
        }
        for future in as_completed(future_map):
            if _INTERRUPTED:
                break
            idx = future_map[future]
            try:
                results[idx] = future.result()
            except Exception as exc:  # pylint: disable=broad-except
                results[idx] = {
                    "url": batch[idx],
                    "cdn_url": "",
                    "error_msg": str(exc),
                }
            completed += 1
            if progress_every > 0 and completed % progress_every == 0:
                done = completed_offset + completed
                print(f"progress: {done}/{total}", file=sys.stderr)
    except KeyboardInterrupt:
        _INTERRUPTED = True
    finally:
        executor.shutdown(wait=False, cancel_futures=True)

    if _INTERRUPTED:
        for idx, item in enumerate(results):
            if item is None:
                results[idx] = {
                    "url": batch[idx],
                    "cdn_url": "",
                    "error_msg": "interrupted",
                }

    return [item for item in results if item is not None]


def build_csv_output_rows(
    rows: Sequence[dict],
    headers: Sequence[str],
    results: Sequence[dict],
) -> List[dict]:
    output: List[dict] = []
    for row, result in zip(rows, results):
        record = {name: row.get(name, "") for name in headers}
        record["CDNURL"] = result.get("cdn_url", "")
        record["error_msg"] = result.get("error_msg", "")
        output.append(record)
    return output


def main() -> int:
    signal.signal(signal.SIGINT, _handle_sigint)
    parser = argparse.ArgumentParser(
        description="Resolve Kuaishou share links to CDN URLs and output JSONL."
    )
    parser.add_argument("input", nargs="?", help="Share text or URL")
    parser.add_argument("--input-csv", help="CSV file with URL field and other columns")
    parser.add_argument(
        "--csv-url-field",
        default="URL",
        help="CSV column name containing the share URL",
    )
    parser.add_argument("--output", help="Write JSONL to this file (default: stdout)")
    parser.add_argument(
        "--resume",
        action="store_true",
        help="Skip URLs already resolved in output JSONL (CSV mode only)",
    )
    parser.add_argument(
        "--proxy",
        default=None,
        help="Proxy URL, e.g. http://user:pass@host:port",
    )
    parser.add_argument("--cookie", default=None, help="Raw Cookie header value")
    parser.add_argument(
        "--cookie-file",
        default=None,
        help="Path to a cookie.txt file or JSON cookie array",
    )
    parser.add_argument("--workers", type=int, default=5, help="Concurrent workers")
    parser.add_argument("--timeout", type=float, default=15.0, help="HTTP timeout (s)")
    parser.add_argument(
        "--batch-size", type=int, default=10, help="Batch size for append output"
    )
    parser.add_argument(
        "--progress-every",
        type=int,
        default=1,
        help="Print progress every N items (stderr)",
    )
    args = parser.parse_args()

    if not args.input and not args.input_csv:
        parser.error("Provide input text/URL or --input-csv")
    if args.input_csv and args.input:
        parser.error("Use only one input source: --input-csv or text")
    if args.resume and not args.input_csv:
        parser.error("--resume is only supported with --input-csv")

    anonymous_cookie = None
    if not args.cookie and not args.cookie_file:
        anonymous_cookie = fetch_anonymous_cookie(args.timeout, args.proxy)

    if args.input_csv:
        if args.resume and (not args.output or args.output == "-"):
            parser.error("--resume requires --output file")
        try:
            rows, headers, url_field = load_csv_rows(
                args.input_csv, args.csv_url_field
            )
        except ValueError as exc:
            parser.error(str(exc))
        if not rows:
            parser.error("No usable CSV rows found")

        skipped = 0
        total = len(rows)
        if args.resume:
            kept, removed = clean_output_jsonl(
                args.output,
                cdn_field="CDNURL",
                error_field="error_msg",
            )
            if removed:
                print(f"resume: removed {removed} failed rows", file=sys.stderr)
            resume_urls = load_resume_success_urls(
                args.output,
                url_field,
                cdn_field="CDNURL",
                error_field="error_msg",
            )
            if resume_urls:
                rows = [
                    row
                    for row in rows
                    if row.get(url_field, "").strip() not in resume_urls
                ]
                skipped = total - len(rows)
                if skipped:
                    print(
                        f"resume: skipped {skipped} already resolved",
                        file=sys.stderr,
                    )
            if not rows:
                print("resume: all rows already resolved", file=sys.stderr)
                return 0

        total_remaining = len(rows)
        batch_size = max(1, args.batch_size)
        completed = skipped
        first_write = not args.resume
        for start in range(0, total_remaining, batch_size):
            batch_rows = rows[start : start + batch_size]
            batch_inputs = [row.get(url_field, "") for row in batch_rows]
            results = process_batch(
                batch=batch_inputs,
                timeout=args.timeout,
                proxy=args.proxy,
                cookie=args.cookie,
                cookie_file=args.cookie_file,
                workers=args.workers,
                progress_every=args.progress_every,
                completed_offset=completed,
                total=total,
                anonymous_cookie=anonymous_cookie,
            )
            output_rows = build_csv_output_rows(batch_rows, headers, results)
            write_jsonl(output_rows, args.output, append=not first_write)
            first_write = False
            completed += len(batch_rows)
            if args.progress_every > 0:
                print(f"progress: {completed}/{total}", file=sys.stderr)
            if _INTERRUPTED:
                print("Interrupted by user, exiting gracefully.", file=sys.stderr)
                return 130
        return 0

    inputs = load_inputs(args.input)
    if not inputs:
        parser.error("No usable input lines found")

    batch_size = max(1, args.batch_size)
    total = len(inputs)
    completed = 0
    first_write = True
    for start in range(0, total, batch_size):
        batch = inputs[start : start + batch_size]
        results = process_batch(
            batch=batch,
            timeout=args.timeout,
            proxy=args.proxy,
            cookie=args.cookie,
            cookie_file=args.cookie_file,
            workers=args.workers,
            progress_every=args.progress_every,
            completed_offset=completed,
            total=total,
            anonymous_cookie=anonymous_cookie,
        )
        write_jsonl(results, args.output, append=not first_write)
        first_write = False
        completed += len(batch)
        if args.progress_every > 0:
            print(f"progress: {completed}/{total}", file=sys.stderr)
        if _INTERRUPTED:
            print("Interrupted by user, exiting gracefully.", file=sys.stderr)
            return 130
    return 0


if __name__ == "__main__":
    raise SystemExit(main())