openapi/source/pixiv/pixiv.py

from flask import Flask, jsonify, send_file
import requests
import os
import json
import urllib3
from pathlib import Path

IMAGE_PROXY_DOMAIN = "openapi.rrrosmontis.icu"
TOKEN_FILE = "tokens.json"
API_HOST = "app-api.pixiv.net"
USER_AGENT = "PixivIOSApp/7.13.0 (iOS 14.6; iPhone)"

urllib3.disable_warnings()

IMAGE_DIR = (Path(__file__).resolve().parent / "../cache/pixiv/images").resolve()
IMAGE_DIR.mkdir(parents=True, exist_ok=True)

app = Flask(__name__)


def get_access_token():
    token_path = Path(__file__).parent / TOKEN_FILE
    if not token_path.exists():
        return None
    try:
        with open(token_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
        return data.get("access_token")
    except Exception:
        return None


def fetch_illust_url(pid, page):
    access_token = get_access_token()
    if not access_token:
        return None
    try:
        resp = requests.get(
            f"https://{API_HOST}/v1/illust/detail",
            params={"illust_id": pid},
            headers={
                "Host": API_HOST,
                "User-Agent": USER_AGENT,
                "Authorization": f"Bearer {access_token}"
            },
            timeout=20,
            verify=False
        )
        if resp.status_code != 200:
            return None
        data = resp.json()
        if data.get("error"):
            return None
        illust = data.get("illust")
        if not illust:
            return None
        page_count = illust.get("page_count", 1)
        if page >= page_count:
            return None
        original_url = None
        if page_count > 1:
            meta_pages = illust.get("meta_pages")
            if meta_pages and page < len(meta_pages):
                original_url = meta_pages[page]["image_urls"]["original"]
        else:
            meta_single = illust.get("meta_single_page")
            if meta_single:
                original_url = meta_single.get("original_image_url")
        if not original_url:
            original_url = illust.get("image_urls", {}).get("large")
        if not original_url:
            return None
        return original_url.replace("i.pximg.net", IMAGE_PROXY_DOMAIN)
    except Exception:
        return None


def make_re_url(pid, page):
    if page == 0:
        return f"https://pixiv.re/{pid}.jpg"
    else:
        return f"https://pixiv.re/{pid}_p{page}.jpg"


def make_cat_url(pid, page):
    if page == 0:
        return f"https://pixiv.cat/{pid}.jpg"
    else:
        return f"https://pixiv.cat/{pid}-{page+1}.jpg"


def check_cache(pid, page):
    pattern = f"{pid}.*" if page == 0 else f"{pid}-p{page+1}.*"
    for f in IMAGE_DIR.glob(pattern):
        ext = f.suffix.lstrip('.')
        mime = 'image/jpeg'
        if ext == 'png':
            mime = 'image/png'
        elif ext == 'gif':
            mime = 'image/gif'
        return f, mime
    return None, None


def download_image(url, pid, page):
    cached_path, mime = check_cache(pid, page)
    if cached_path:
        return cached_path, mime
    try:
        resp = requests.get(
            url,
            headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"},
            timeout=30,
            verify=False
        )
        if resp.status_code != 200:
            return None, None
        ct = resp.headers.get('Content-Type', 'image/jpeg')
        ext = 'jpg'
        if 'png' in ct:
            ext = 'png'
        elif 'gif' in ct:
            ext = 'gif'
        filename = f"{pid}.{ext}" if page == 0 else f"{pid}-p{page+1}.{ext}"
        tmp = IMAGE_DIR / (filename + '.tmp')
        final = IMAGE_DIR / filename
        with open(tmp, 'wb') as f:
            f.write(resp.content)
        os.replace(tmp, final)
        return final, ct
    except Exception:
        return None, None


@app.route('/<path:path>')
def serve_image(path):
    import re
    m = re.match(r'^(\d+)(?:-(\d+))?(?:\.\w+)?$', path)
    if not m:
        return jsonify(error="Invalid path format, expected /{pid} or /{pid}-{page}"), 400
    pid = int(m.group(1))
    page = int(m.group(2)) - 1 if m.group(2) else 0

    cached_file, mime = check_cache(pid, page)
    if cached_file:
        return send_file(cached_file, mimetype=mime, max_age=3600)

    urls_to_try = []
    official_url = fetch_illust_url(pid, page)
    if official_url:
        urls_to_try.append(official_url)
    urls_to_try.append(make_re_url(pid, page))
    urls_to_try.append(make_cat_url(pid, page))

    for url in urls_to_try:
        img_path, mime = download_image(url, pid, page)
        if img_path:
            return send_file(img_path, mimetype=mime, max_age=3600)

    return jsonify(error="Failed to fetch image from all sources"), 500


@app.route('/health')
def health():
    access_token = get_access_token()
    if not access_token:
        return jsonify(status="error", detail="tokens.json missing or invalid")
    try:
        resp = requests.get(
            f"https://{API_HOST}/v1/illust/detail",
            params={"illust_id": 105383871},
            headers={
                "Host": API_HOST,
                "User-Agent": USER_AGENT,
                "Authorization": f"Bearer {access_token}"
            },
            timeout=15,
            verify=False
        )
        if resp.status_code == 200:
            return jsonify(status="ok", token="valid")
        else:
            return jsonify(status="error", detail=f"API returned {resp.status_code}")
    except Exception as e:
        return jsonify(status="error", detail=str(e))


if __name__ == '__main__':
    token_path = Path(__file__).parent / TOKEN_FILE
    if not token_path.exists():
        print(f"Warning: {TOKEN_FILE} not found, access_token may be missing")
    else:
        token = get_access_token()
        if token:
            print("Access token loaded")
        else:
            print("Warning: access_token missing in tokens.json")
    print(f"Cache directory: {IMAGE_DIR}")
    app.run(host='127.0.0.1', port=9007, debug=False, threaded=True)