wos_export.py

#!/usr/bin/env python3
"""Export Web of Science search results from a summary URL."""

from __future__ import annotations

import argparse
import json
import os
import re
import time
import urllib.error
import urllib.parse
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path


FIELDS = [
    "AUTHORS",
    "TITLE",
    "SOURCE",
    "DOI",
    "CONFERENCE_INFO",
    "CITTIMES",
    "ACCESSION_NUM",
    "AUTHORSIDENTIFIERS",
    "ISSN",
    "PMID",
    "ABSTRACT",
    "ADDRS",
    "AFFILIATIONS",
    "DOCTYPE",
    "KEYWORDS",
    "JCR_CATEGORY",
    "SUBJECT_CATEGORY",
    "WOS_EDITIONS",
    "CITREF",
    "CITREFC",
    "USAGEIND",
    "HOT_PAPER",
    "HIGHLY_CITED",
    "FUNDING",
    "PUBINFO",
    "OPEN_ACCESS",
    "PAGEC",
    "SABBR",
    "IDS",
    "LANG",
]

UUID_LIKE_ID_RE = re.compile(
    r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}(?:-[0-9a-f]+)?$",
    re.IGNORECASE,
)
STORAGE_SID_RE = re.compile(rb"\b[A-Za-z0-9][A-Za-z0-9._~=-]{15,80}\b")


def build_referer(origin: str, parent_qid: str, sid: str) -> str:
    if sid:
        return f"{origin}/wos/woscc/summary/{parent_qid}/{sid}/relevance/1(overlay:export/exp)"
    return f"{origin}/wos/woscc/summary/{parent_qid}/relevance/1(overlay:export/exp)"


def is_likely_export_sid(value: str) -> bool:
    value = value.strip()
    if not value or "/" in value or "\\" in value or any(ch.isspace() for ch in value):
        return False
    if UUID_LIKE_ID_RE.fullmatch(value):
        return False
    return bool(re.fullmatch(r"[A-Za-z0-9._~=-]{8,}", value))


def is_likely_storage_sid(value: str) -> bool:
    if not is_likely_export_sid(value):
        return False
    if not (20 <= len(value) <= 64):
        return False
    if not re.fullmatch(r"[A-Za-z0-9._~=-]+", value):
        return False
    lowered = value.lower()
    if lowered.startswith(("wos_", "snowplow", "current", "notice", "stats", "tabmanager")):
        return False
    return any(ch.isalpha() for ch in value) and any(ch.isdigit() for ch in value)


def parse_summary_input(value: str) -> dict[str, str]:
    value = value.strip()
    if not value:
        raise ValueError("empty QID or URL")

    if "://" not in value:
        origin = "https://www.webofscience.com"
        return {
            "parent_qid": value,
            "sid": "",
            "route_id": "",
            "origin": origin,
            "referer": build_referer(origin, value, ""),
        }

    parsed = urllib.parse.urlparse(value)
    if not parsed.scheme or not parsed.netloc:
        raise ValueError("invalid Web of Science URL")

    parts = [part for part in parsed.path.split("/") if part]
    try:
        index = parts.index("summary")
    except ValueError as exc:
        raise ValueError("could not find '/summary/' in URL") from exc

    if index + 1 >= len(parts):
        raise ValueError("could not find QID after '/summary/'")

    parent_qid = parts[index + 1]
    route_id = ""
    if index + 2 < len(parts) and parts[index + 2] != "relevance":
        route_id = parts[index + 2]
    sid = route_id if is_likely_export_sid(route_id) else ""

    origin = f"{parsed.scheme}://{parsed.netloc}"
    return {
        "parent_qid": parent_qid,
        "sid": sid,
        "route_id": route_id,
        "origin": origin,
        "referer": build_referer(origin, parent_qid, route_id),
    }


def extract_sid_from_url(url: str) -> str:
    parsed = urllib.parse.urlparse(url)
    query = urllib.parse.parse_qs(parsed.query)
    referrer = query.get("referrer", [""])[0]
    match = re.search(r"SID=([^&]+)", referrer)
    if match:
        return match.group(1)
    return query.get("SID", [""])[0]


def browser_user_data_roots() -> list[Path]:
    local_app_data = Path(os.environ.get("LOCALAPPDATA", ""))
    if not local_app_data:
        return []
    return [
        local_app_data / "Google" / "Chrome" / "User Data",
        local_app_data / "Microsoft" / "Edge" / "User Data",
        local_app_data / "BraveSoftware" / "Brave-Browser" / "User Data",
    ]


def browser_storage_leveldb_dirs() -> list[Path]:
    dirs = []
    for root in browser_user_data_roots():
        if not root.exists():
            continue
        for profile in root.iterdir():
            leveldb = profile / "Local Storage" / "leveldb"
            if leveldb.is_dir():
                dirs.append(leveldb)
    return dirs


def sid_storage_needles(origin: str) -> list[bytes]:
    origins = {
        "https://webofscience.clarivate.cn",
        "https://www.webofscience.com",
        "https://www.webofknowledge.com",
    }
    parsed = urllib.parse.urlparse(origin)
    if parsed.scheme and parsed.netloc:
        origins.add(f"{parsed.scheme}://{parsed.netloc}")

    needles = set()
    for item in origins:
        parsed_item = urllib.parse.urlparse(item)
        needles.add(item.lower().encode("utf-8"))
        if parsed_item.netloc:
            needles.add(parsed_item.netloc.lower().encode("utf-8"))
    return sorted(needles)


def iter_browser_storage_sid_candidates(origin: str) -> list[str]:
    needles = sid_storage_needles(origin)
    candidates: dict[str, float] = {}
    for leveldb in browser_storage_leveldb_dirs():
        for path in list(leveldb.glob("*.ldb")) + list(leveldb.glob("*.log")):
            try:
                if path.stat().st_size > 50_000_000:
                    continue
                data = path.read_bytes()
            except OSError:
                continue

            lowered = data.lower()
            if not any(needle in lowered for needle in needles):
                continue

            interesting_offsets: list[int] = []
            for needle in [*needles, b"sid"]:
                start = 0
                while len(interesting_offsets) < 300:
                    index = lowered.find(needle, start)
                    if index < 0:
                        break
                    interesting_offsets.append(index)
                    start = index + 1

            modified_at = path.stat().st_mtime
            for index in interesting_offsets:
                snippet = data[max(0, index - 300) : index + 800]
                snippet_lowered = snippet.lower()
                if b"sid" not in snippet_lowered and not any(needle in snippet_lowered for needle in needles):
                    continue
                for match in STORAGE_SID_RE.finditer(snippet):
                    try:
                        candidate = match.group(0).decode("ascii")
                    except UnicodeDecodeError:
                        continue
                    if is_likely_storage_sid(candidate):
                        candidates[candidate] = max(candidates.get(candidate, 0), modified_at)

    def priority(item: tuple[str, float]) -> tuple[int, float, int, str]:
        candidate, modified_at = item
        has_upper = any(ch.isupper() for ch in candidate)
        has_lower = any(ch.islower() for ch in candidate)
        has_digit = any(ch.isdigit() for ch in candidate)
        rank = 100
        if candidate.isalnum() and 24 <= len(candidate) <= 40 and has_upper and has_lower and has_digit:
            rank -= 60
        if candidate.startswith(("USW", "USE", "EUW", "EUC", "AP")):
            rank -= 30
        if not candidate.isalnum():
            rank += 20
        if len(candidate) > 48:
            rank += 20
        if candidate.lower().startswith(("wos", "snow", "queue", "track", "suite")):
            rank += 40
        return (rank, -modified_at, abs(len(candidate) - 30), candidate)

    return [candidate for candidate, _ in sorted(candidates.items(), key=priority)][:80]


def validate_sid(origin: str, sid: str) -> bool:
    if not sid:
        return False
    origin = origin.strip().rstrip("/") or "https://webofscience.clarivate.cn"
    body = b"{}"
    headers = {
        "User-Agent": (
            "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
            "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126 Safari/537.36"
        ),
        "Accept": "application/json",
        "Content-Type": "application/json",
        "Origin": origin,
        "Referer": f"{origin}/wos/",
        "x-1p-wos-sid": sid,
    }
    request = urllib.request.Request(
        f"{origin}/api/esti/Session/getSessionData",
        data=body,
        headers=headers,
        method="POST",
    )
    try:
        with urllib.request.urlopen(request, timeout=12) as response:
            payload = json.loads(response.read().decode("utf-8", errors="replace"))
    except (urllib.error.URLError, ValueError, json.JSONDecodeError):
        return False
    return isinstance(payload, dict) and (
        "BasicProperties" in payload or "Products" in payload or "CustomerProperties" in payload
    )


def get_sid_from_browser_storage(origin: str) -> str:
    for candidate in iter_browser_storage_sid_candidates(origin):
        if validate_sid(origin, candidate):
            return candidate
    return ""


def get_sid(origin: str) -> str:
    candidates = []
    origin = origin.strip().rstrip("/")
    sid = get_sid_from_browser_storage(origin)
    if sid:
        return sid

    if origin:
        candidates.append(f"{origin}/")
    candidates.extend(
        [
            "https://webofscience.clarivate.cn/",
            "https://www.webofscience.com/",
            "https://www.webofknowledge.com/",
            "http://www.webofknowledge.com/",
        ]
    )

    seen = set()
    headers = {
        "User-Agent": (
            "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
            "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/126 Safari/537.36"
        ),
        "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
        "Accept-Language": "zh-CN,zh;q=0.9,en-US;q=0.7,en;q=0.3",
    }

    for candidate in candidates:
        if candidate in seen:
            continue
        seen.add(candidate)
        try:
            request = urllib.request.Request(candidate, headers=headers)
            response = urllib.request.urlopen(request, timeout=30)
            sid = extract_sid_from_url(response.geturl())
            response.close()
        except (urllib.error.URLError, ValueError):
            continue
        if sid:
            return sid
    return ""


def request_range(ctx: dict[str, str], sid: str, start: int, end: int, retries: int = 3) -> bytes:
    payload = {
        "parentQid": ctx["parent_qid"],
        "sortBy": "relevance",
        "displayTimesCited": "true",
        "displayCitedRefs": "true",
        "product": "UA",
        "colName": "WOS",
        "displayUsageInfo": "true",
        "fileOpt": "othersoftware",
        "action": "saveToFieldTagged",
        "markFrom": str(start),
        "markTo": str(end),
        "view": "summary",
        "isRefQuery": "false",
        "locale": "en_US",
        "fieldList": FIELDS,
    }
    body = json.dumps(payload).encode("utf-8")
    headers = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; rv:109.0) Gecko/20100101 Firefox/119.0",
        "Accept": "application/json, text/plain, */*",
        "Accept-Language": "zh-CN,zh;q=0.9,en-US;q=0.7,en;q=0.3",
        "X-1P-WOS-SID": sid,
        "Content-Type": "application/json",
        "Origin": ctx["origin"],
        "DNT": "1",
        "Connection": "keep-alive",
        "Referer": ctx["referer"],
        "Sec-Fetch-Dest": "empty",
        "Sec-Fetch-Mode": "cors",
        "Sec-Fetch-Site": "same-origin",
    }
    url = f"{ctx['origin']}/api/wosnx/indic/export/saveToFile"

    for attempt in range(1, retries + 1):
        request = urllib.request.Request(url, data=body, headers=headers, method="POST")
        try:
            with urllib.request.urlopen(request, timeout=90) as response:
                return response.read()
        except urllib.error.HTTPError as exc:
            if exc.code == 504 and attempt < retries:
                time.sleep(2 * attempt)
                continue
            detail = exc.read().decode("utf-8", errors="replace")
            raise RuntimeError(f"HTTP {exc.code} for {start}-{end}: {detail[:500]}") from exc
        except urllib.error.URLError as exc:
            if attempt < retries:
                time.sleep(2 * attempt)
                continue
            raise RuntimeError(f"Request failed for {start}-{end}: {exc}") from exc

    raise RuntimeError(f"Request failed for {start}-{end}")


def export_range(ctx: dict[str, str], sid: str, output_dir: Path, start: int, end: int, overwrite: bool) -> Path:
    output_path = output_dir / f"{start:06d}_{end:06d}.txt"
    if output_path.exists() and output_path.stat().st_size > 0 and not overwrite:
        print(f"Skipping existing file: {output_path}")
        return output_path

    data = request_range(ctx, sid, start, end)
    if not data:
        raise RuntimeError(f"Empty response for {start}-{end}")
    output_path.write_bytes(data)
    print(f"Saved: {output_path}")
    return output_path


def ranges(total: int) -> list[tuple[int, int]]:
    return [(start, min(start + 999, total)) for start in range(1, total + 1, 1000)]


def sanitize_folder_name(value: str) -> str:
    value = re.sub(r"[<>:\"/\\|?*\x00-\x1f]+", "_", value.strip())
    value = re.sub(r"\s+", "_", value)
    return value.strip(" ._")[:120]


def default_output_dir(parent_qid: str) -> Path:
    stamp = time.strftime("%Y%m%d_%H%M%S")
    short_qid = sanitize_folder_name(parent_qid[:8] or "wos")
    return Path(__file__).resolve().parent / "exports" / f"wos_{stamp}_{short_qid}"


def main() -> int:
    parser = argparse.ArgumentParser(description="Export Web of Science results in 1000-record chunks.")
    parser.add_argument("summary", help="Full Web of Science summary URL, or the old single QID.")
    parser.add_argument("total", type=int, help="Total number of records to export.")
    parser.add_argument("-o", "--output-dir", help="Output directory. Default: the parsed parent QID.")
    parser.add_argument("--sid", default="", help="Optional SID override for single-QID URLs.")
    parser.add_argument("--workers", type=int, default=6, help="Concurrent export requests. Default: 6.")
    parser.add_argument("--overwrite", action="store_true", help="Overwrite existing split files.")
    args = parser.parse_args()

    ctx = parse_summary_input(args.summary)
    sid = args.sid.strip() or ctx["sid"] or get_sid(ctx["origin"])
    if not sid:
        raise SystemExit(
            "SID not found. Single-QID mode requires the tool to auto-fetch a SID from Web of Knowledge, "
            "but this environment did not return one. Paste a URL like "
            "/wos/woscc/summary/<QID>/<SID>/relevance/1 or pass --sid manually."
        )

    output_dir = Path(args.output_dir).resolve() if args.output_dir else default_output_dir(ctx["parent_qid"]).resolve()
    output_dir.mkdir(parents=True, exist_ok=True)

    print(f"Origin: {ctx['origin']}")
    print(f"Parent QID: {ctx['parent_qid']}")
    print(f"SID: {sid}")
    print(f"Output directory: {output_dir}")

    chunks = ranges(args.total)
    with ThreadPoolExecutor(max_workers=args.workers) as executor:
        futures = [
            executor.submit(export_range, ctx, sid, output_dir, start, end, args.overwrite)
            for start, end in chunks
        ]
        for future in as_completed(futures):
            future.result()

    print("Export complete.")
    return 0


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