import asyncio import base64 import hashlib import io import json import os import re import secrets import tempfile import time import zipfile import zlib from pathlib import Path from urllib.parse import parse_qs, parse_qsl, quote, unquote, urlencode, urlsplit try: from dotenv import load_dotenv load_dotenv(Path(__file__).resolve().parent.parent / ".env") except ImportError: pass import aiohttp import httpx import numpy as np from fastapi import BackgroundTasks, FastAPI, HTTPException, Request from fastapi.responses import FileResponse, HTMLResponse, Response, StreamingResponse from fastapi.staticfiles import StaticFiles from cryptography.hazmat.primitives.ciphers.aead import AESGCM from PIL import ExifTags, Image from pydantic import BaseModel TURSO_DB_URL = os.getenv("TURSO_DB_URL", "").strip() TURSO_AUTH_TOKEN = os.getenv("TURSO_AUTH_TOKEN_WRITE", "").strip() PHPSESSID = os.getenv("PHPSESSID", "") IMG_BASE = "https://i.pximg.net/img-original/img/" PIXIV_HEADERS = { "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:126.0) Gecko/20100101 Firefox/126.0", "referer": "https://www.pixiv.net/", } AI_TAGS = { "ai", "ai-generated", "aiart", "aiillustration", "artai", "comfyui", "novelai", "novelaidiffusionai", "pixai", "stablediffusion", "aiイラスト", "イラストai", "ai生成", "生成ai", "ai作成", "ai绘画", } EXIF_TYPE_ORDER = ( "nai", "comfy", "sd", "mj", "celsys", "photoshop", "nai_stealth", "invoke", "fooocus", "swarm", "nai_vibe", ) EXIF_TYPE_TO_CODE = {name: idx + 1 for idx, name in enumerate(EXIF_TYPE_ORDER)} EXIF_CODE_TO_TYPE = {idx + 1: name for idx, name in enumerate(EXIF_TYPE_ORDER)} FALLBACK_EXIF_TYPES = {"celsys", "photoshop"} POST_SCAN_LIMIT = 64 FULL_IMAGE_LIMIT = 32 POPULARITY_BATCH_SIZE = 32 DOWNLOAD_BATCH_SIZE = 8 PAGE_SIZE = 60 SEARCH_PAGE_SIZE = 10 RECLASS_BATCH_SIZE = 64 REFRESH_SEARCH_LIMIT = 128 SEARCH_LIST_WHERE = "post_ids != '[]' AND found_exif > 0" CACHE_EXPLORER_PAGE_1 = "explorer:page:1" CACHE_EXPLORER_TOTAL = "explorer:total" POST_SEARCH_LIMIT = 1000 RECLASS_LOG_KEY = "TF#g9FA@UgKdN'usVtm3dTz*yoqbria`" IMAGE_CACHE_DIR = Path("/data") IMAGE_CACHE_PASSWORD = '68c6I5B!W[6/"0n>W%G/lUOEY`K|zvQ4' IMAGE_CACHE_KEY = hashlib.sha256(IMAGE_CACHE_PASSWORD.encode()).digest() IMAGE_URL_VERSION = 5 THUMB_SIZE = 540 THUMB_MAX_AGE = 1800 THUMB_CLEAN_INTERVAL = 300 THUMB_DIR = Path(tempfile.gettempdir()) / "pixif2-thumbs" PAGE_URL_CACHE_MAX_AGE = 1800 PAGE_URL_CACHE_LIMIT = 5000 WEBP_SCALE = 0.4 WEBP_QUALITY = 82 STEALTH_ALPHA_SIGS = (b"stealth_pnginfo", b"stealth_pngcomp") STEALTH_RGB_SIGS = (b"stealth_rgbinfo", b"stealth_rgbcomp") EXIF_FIELD_NAMES = { "DocumentName": "Title", "ImageDescription": "Description", "UserComment": "Comment", } app = FastAPI() ACTIVE_TASKS = {} TASK_CANCELS = set() TASK_SCAN_TASKS = {} TASK_EVENT_QUEUES = set() PAGE_URL_CACHE = {} PAGE_URL_INFLIGHT = {} IMAGE_CACHE_LOCKS = [asyncio.Lock() for _ in range(FULL_IMAGE_LIMIT)] POST_SCAN_SEM = asyncio.Semaphore(POST_SCAN_LIMIT) FULL_IMAGE_SEM = asyncio.Semaphore(FULL_IMAGE_LIMIT) TURSO_CLIENT = None PIXIV_CLIENT = None LAST_THUMB_CLEANUP = 0 FRONTEND_DIR = os.path.join(os.getcwd(), "frontend") class TaskCanceled(Exception): pass def task_canceled(task_id): return bool( task_id and (task_id in TASK_CANCELS or ACTIVE_TASKS.get(task_id, {}).get("cancel")) ) def start_task(task_id, data): if task_canceled(task_id): return False ACTIVE_TASKS[task_id] = data return True def cancel_task(task_id): if not task_id: return TASK_CANCELS.add(task_id) active = ACTIVE_TASKS.get(task_id) if active: active["cancel"] = True active["phase"] = "canceling" for task in TASK_SCAN_TASKS.get(task_id, ()): task.cancel() def base26_time(): x = "" n = int(time.time() * 100) while n: x = chr(97 + n % 26) + x n //= 26 return x def base26_to_time(value): n = 0 for c in value.split("_", 1)[0]: if c < "a" or c > "z": return 0 n = n * 26 + ord(c) - 97 return n // 100 def turso_url(): base = TURSO_DB_URL.rstrip("/") if not base.startswith("http"): base = "https://" + base return base async def turso_client(): global TURSO_CLIENT if TURSO_CLIENT is None or TURSO_CLIENT.is_closed: TURSO_CLIENT = httpx.AsyncClient(timeout=120) return TURSO_CLIENT async def pixiv_client(): global PIXIV_CLIENT if PIXIV_CLIENT is None or PIXIV_CLIENT.closed: PIXIV_CLIENT = aiohttp.ClientSession( cookies={"PHPSESSID": PHPSESSID}, headers=PIXIV_HEADERS ) return PIXIV_CLIENT async def turso_execute(stmts): url = turso_url() + "/v2/pipeline" headers = { "Authorization": f"Bearer {TURSO_AUTH_TOKEN}", "Content-Type": "application/json", } body = { "requests": [{"type": "execute", "stmt": s} for s in stmts] + [{"type": "close"}] } client = await turso_client() r = await client.post(url, json=body, headers=headers) r.raise_for_status() return r.json() async def turso_batch(stmts): return await turso_execute(stmts) async def init_db(): await turso_execute( [ { "sql": "CREATE TABLE IF NOT EXISTS pi_searches (id TEXT PRIMARY KEY, post_ids TEXT)" }, { "sql": "CREATE TABLE IF NOT EXISTS pi_scans (post_id TEXT PRIMARY KEY, url TEXT, exif_type INTEGER)" }, { "sql": "CREATE TABLE IF NOT EXISTS pi_scan_pages (" "post_id TEXT, page INTEGER, url TEXT, exif_type INTEGER, " "PRIMARY KEY (post_id, page))" }, {"sql": "CREATE TABLE IF NOT EXISTS pi_meta (key TEXT PRIMARY KEY, value TEXT)"}, { "sql": "CREATE TABLE IF NOT EXISTS pi_cache (key TEXT PRIMARY KEY, value TEXT, updated_at INTEGER)" }, { "sql": "CREATE TABLE IF NOT EXISTS pi_reclass_jobs (" "id TEXT PRIMARY KEY, post_ids TEXT, filters TEXT, total INTEGER DEFAULT 0, " "scanned INTEGER DEFAULT 0, changed INTEGER DEFAULT 0, summary TEXT, " "log_enc TEXT, created_at INTEGER)" }, ] ) for sql in ( "DROP INDEX IF EXISTS pi_searches_created_at_idx", "DROP INDEX IF EXISTS pi_search_posts_search_pos_idx", "DROP INDEX IF EXISTS pi_search_posts_post_idx", "DROP TABLE IF EXISTS pi_search_posts", "ALTER TABLE pi_searches DROP COLUMN api_url", "ALTER TABLE pi_searches DROP COLUMN created_at", "ALTER TABLE pi_searches ADD COLUMN found_exif INTEGER DEFAULT 0", "ALTER TABLE pi_searches ADD COLUMN total_searched INTEGER DEFAULT 0", "ALTER TABLE pi_searches ADD COLUMN deep_scan INTEGER DEFAULT 1", "ALTER TABLE pi_searches ADD COLUMN deep_scan_types TEXT DEFAULT '[]'", "ALTER TABLE pi_scans ADD COLUMN deep_scanned INTEGER DEFAULT 0", "ALTER TABLE pi_scans ADD COLUMN bookmark_count INTEGER", "ALTER TABLE pi_scans ADD COLUMN view_count INTEGER", ): try: await turso_execute([{"sql": sql}]) except httpx.HTTPStatusError as e: text = e.response.text.casefold() if "no such column" not in text and "no such index" not in text and "duplicate column" not in text: raise await turso_execute( [ { "sql": "CREATE INDEX IF NOT EXISTS pi_searches_refresh_idx " "ON pi_searches(total_searched, id)" } ] ) await migrate_exif_order() async def migrate_exif_order(): order_value = ",".join(EXIF_TYPE_ORDER) previous_order = "novelai,comfy,sd,mj,celsys,photoshop,stealth,invoke,fooocus,swarm" swapped_order = "novelai,comfy,sd,mj,celsys,photoshop,stealth" legacy_order = "novelai,sd,comfy,mj,celsys,photoshop,stealth" resp = await turso_execute( [ { "sql": "SELECT value FROM pi_meta WHERE key = 'exif_order'", } ] ) rows = first_rows(resp) prev_order = cell_value(rows[0], 0, "") if rows else "" if prev_order == order_value: return stmts = [] if prev_order == legacy_order: stmts.extend( [ { "sql": "UPDATE pi_scans SET exif_type = CASE exif_type " "WHEN 2 THEN -2 WHEN 3 THEN -3 ELSE exif_type END" }, {"sql": "UPDATE pi_scans SET exif_type = 3 WHERE exif_type = -2"}, {"sql": "UPDATE pi_scans SET exif_type = 2 WHERE exif_type = -3"}, ] ) if prev_order in ("", previous_order, swapped_order, legacy_order): stmts.append( { "sql": "INSERT OR REPLACE INTO pi_meta " "(key, value) VALUES ('exif_order', ?)", "args": [{"type": "text", "value": order_value}], } ) if stmts: await turso_execute(stmts) async def publish_task_event(search_id): data = {"id": search_id, "at": int(time.time())} for queue in list(TASK_EVENT_QUEUES): if queue.full(): try: queue.get_nowait() except asyncio.QueueEmpty: pass queue.put_nowait(data) async def finish_task(search_id): ACTIVE_TASKS.pop(search_id, None) TASK_SCAN_TASKS.pop(search_id, None) await publish_task_event(search_id) def is_ai_post(post): if post.get("aiType") == 2: return True tags = post.get("tags") or [] for tag in tags: name = ( tag if isinstance(tag, str) else (tag.get("tag") or tag.get("name") or "") if isinstance(tag, dict) else "" ) if name and name.casefold() in AI_TAGS: return True return False def get_search_keywords(raw): parts = urlsplit(raw) path_parts = [unquote(p) for p in parts.path.split("/") if p] if "tags" in path_parts: idx = path_parts.index("tags") + 1 if idx < len(path_parts): return path_parts[idx] query = parse_qs(parts.query) words = query.get("word") or query.get("q") if words: return words[0] return raw.strip() def get_search_params(raw, keywords): params = [] for key, value in parse_qsl(urlsplit(raw).query, keep_blank_values=True): if key in ("p", "q", "word", "type"): continue if key == "s_mode": if value == "tag": value = "s_tag" elif value == "tag_full": value = "s_tag_full" params.append((key, value)) params.append(("word", keywords)) if not any(k == "s_mode" for k, _ in params): params.append(("s_mode", "s_tag")) return urlencode(params) def get_search_api_url(raw, keywords): encoded = quote(keywords, safe="") params = get_search_params(raw, keywords) return f"https://www.pixiv.net/ajax/search/artworks/{encoded}?{params}" def exif_codes(values): return sorted( { int(code) for code in values or [] if str(code).isdigit() and int(code) in EXIF_CODE_TO_TYPE } ) async def save_search(search_id, post_ids, deep_scan=True, deep_scan_types=None): if not post_ids: return if task_canceled(search_id): raise TaskCanceled() deep_scan_types = exif_codes(deep_scan_types) if deep_scan else [] stmt = { "sql": "INSERT OR REPLACE INTO pi_searches " "(id, post_ids, deep_scan, deep_scan_types) VALUES (?, ?, ?, ?)", "args": [ {"type": "text", "value": search_id}, {"type": "text", "value": json.dumps(post_ids)}, {"type": "integer", "value": "1" if deep_scan else "0"}, {"type": "text", "value": json.dumps(deep_scan_types or [])}, ], } await turso_execute([stmt]) if task_canceled(search_id): await delete_task_rows(search_id) raise TaskCanceled() async def delete_task_rows(search_id): await turso_execute( [ { "sql": "DELETE FROM pi_searches WHERE id = ?", "args": [{"type": "text", "value": search_id}], }, { "sql": "DELETE FROM pi_reclass_jobs WHERE id = ?", "args": [{"type": "text", "value": search_id}], }, ] ) def cell_value(row, index, default=None): if index >= len(row) or row[index].get("type") == "null": return default return row[index].get("value", default) def first_rows(resp, index=0): results = resp.get("results") or [] if len(results) <= index or "response" not in results[index]: return [] return results[index]["response"].get("result", {}).get("rows", []) async def set_cache(key, value): await turso_execute( [ { "sql": "INSERT OR REPLACE INTO pi_cache (key, value, updated_at) VALUES (?, ?, ?)", "args": [ {"type": "text", "value": key}, {"type": "text", "value": json.dumps(value)}, {"type": "integer", "value": str(int(time.time()))}, ], } ] ) async def get_cache(key): resp = await turso_execute( [ { "sql": "SELECT value FROM pi_cache WHERE key = ?", "args": [{"type": "text", "value": key}], } ] ) rows = first_rows(resp) if not rows: return None return json.loads(cell_value(rows[0], 0, "null")) def explorer_item(row): search_id = cell_value(row, 0, "") kind = cell_value(row, 1, "search") changed = cell_value(row, 4) total_count = cell_value(row, 5) created_at = int(cell_value(row, 6, "0") or "0") return { "id": search_id, "kind": kind, "created_at": created_at or base26_to_time(search_id), "found_exif": int(cell_value(row, 2, "0")), "total_searched": int(cell_value(row, 3, "0")), "changed": int(changed) if changed else None, "total": int(total_count) if total_count else None, } def explorer_response(items, total, page): return { "items": items, "total": total, "page": page, "page_size": SEARCH_PAGE_SIZE, "pages": max((total + SEARCH_PAGE_SIZE - 1) // SEARCH_PAGE_SIZE, 1), } async def query_explorer_total(): resp = await turso_execute( [ { "sql": "WITH items AS (" f"SELECT id FROM pi_searches WHERE {SEARCH_LIST_WHERE} " "UNION ALL " "SELECT id FROM pi_reclass_jobs" ") SELECT COUNT(*) FROM items" } ] ) rows = first_rows(resp) return int(cell_value(rows[0], 0, "0")) if rows else 0 async def query_explorer_page(page): offset = (page - 1) * SEARCH_PAGE_SIZE item_sql = ( "WITH items AS (" "SELECT id, 'search' kind, found_exif, total_searched, NULL changed, NULL total, 0 created_at " f"FROM pi_searches WHERE {SEARCH_LIST_WHERE} " "UNION ALL " "SELECT id, 'reclass' kind, changed found_exif, scanned total_searched, changed, total, created_at " "FROM pi_reclass_jobs" ") " ) resp = await turso_execute( [ {"sql": item_sql + "SELECT COUNT(*) FROM items"}, { "sql": item_sql + "SELECT id, kind, found_exif, total_searched, changed, total, created_at FROM items " "ORDER BY id DESC LIMIT ? OFFSET ?", "args": [ {"type": "integer", "value": str(SEARCH_PAGE_SIZE)}, {"type": "integer", "value": str(offset)}, ], }, ] ) count_rows = first_rows(resp) total = int(cell_value(count_rows[0], 0, "0")) if count_rows else 0 return explorer_response([explorer_item(row) for row in first_rows(resp, 1)], total, page) async def query_explorer_item(search_id): resp = await turso_execute( [ { "sql": "SELECT id, 'search' kind, found_exif, total_searched, NULL changed, NULL total, 0 created_at " f"FROM pi_searches WHERE id = ? AND {SEARCH_LIST_WHERE} " "UNION ALL " "SELECT id, 'reclass' kind, changed found_exif, scanned total_searched, changed, total, created_at " "FROM pi_reclass_jobs WHERE id = ?", "args": [ {"type": "text", "value": search_id}, {"type": "text", "value": search_id}, ], } ] ) rows = first_rows(resp) return explorer_item(rows[0]) if rows else None async def rebuild_explorer_cache(): data = await query_explorer_page(1) await set_cache(CACHE_EXPLORER_PAGE_1, {"items": data["items"]}) await set_cache(CACHE_EXPLORER_TOTAL, data["total"]) return data async def cached_explorer_page(): page = await get_cache(CACHE_EXPLORER_PAGE_1) total = await get_cache(CACHE_EXPLORER_TOTAL) if page is None or total is None: return None return explorer_response(page.get("items", []), int(total), 1) async def patch_explorer_cache(search_id, front=False): page = await get_cache(CACHE_EXPLORER_PAGE_1) total = await get_cache(CACHE_EXPLORER_TOTAL) if page is None or total is None: await rebuild_explorer_cache() return items = page.get("items", []) index = next((i for i, item in enumerate(items) if item["id"] == search_id), -1) item = await query_explorer_item(search_id) if item is None: if index >= 0: items.pop(index) total = max(int(total) - 1, 0) else: return elif index >= 0: items[index] = item elif front: items.insert(0, item) total = int(total) + 1 else: await set_cache(CACHE_EXPLORER_TOTAL, await query_explorer_total()) return await set_cache(CACHE_EXPLORER_PAGE_1, {"items": items[:SEARCH_PAGE_SIZE]}) await set_cache(CACHE_EXPLORER_TOTAL, int(total)) async def get_search_post_ids(search_id): resp = await turso_execute( [ { "sql": "SELECT post_ids FROM pi_searches WHERE id = ?", "args": [{"type": "text", "value": search_id}], } ] ) rows = first_rows(resp) if not rows: return None return json.loads(cell_value(rows[0], 0, "[]")) async def get_search_counts(search_id): resp = await turso_execute( [ { "sql": RESULT_BASE_SQL + "SELECT COUNT(*), COUNT(s.post_id), " "(SELECT COUNT(DISTINCT post_id) FROM tags) " "FROM ids LEFT JOIN pi_scans s ON s.post_id = ids.post_id", "args": [{"type": "text", "value": search_id}], } ] ) rows = first_rows(resp) if not rows: return None row = rows[0] return { "raw_total": int(cell_value(row, 0, "0")), "scanned_count": int(cell_value(row, 1, "0")), "found_exif": int(cell_value(row, 2, "0")), } async def scan_progress_counts(search_id, batch_total): counts = await get_search_counts(search_id) if not counts: return {"total": batch_total, "done": 0} return { "total": batch_total, "done": 0, "raw_total": counts["raw_total"], "scanned_base": counts["scanned_count"], } async def update_search_stats(search_id): counts = await get_search_counts(search_id) if counts is None: return await turso_execute( [ { "sql": "UPDATE pi_searches SET found_exif = ?, total_searched = ? WHERE id = ?", "args": [ {"type": "integer", "value": str(counts["found_exif"])}, {"type": "integer", "value": str(counts["scanned_count"])}, {"type": "text", "value": search_id}, ], } ] ) async def update_search_stats_batch(search_ids): search_ids = [sid for sid in search_ids if sid] if not search_ids: return 0 updated = 0 for i in range(0, len(search_ids), REFRESH_SEARCH_LIMIT): batch = search_ids[i : i + REFRESH_SEARCH_LIMIT] marks = ",".join("?" for _ in batch) await turso_execute( [ { "sql": "WITH counts AS (" "SELECT p.id, COUNT(s.post_id) scanned, COUNT(CASE WHEN " "s.exif_type IS NOT NULL AND (COALESCE(p.deep_scan_types, '[]') = '[]' " "OR EXISTS (SELECT 1 FROM json_each(p.deep_scan_types) " "WHERE CAST(value AS INTEGER) = s.exif_type) OR EXISTS (" "SELECT 1 FROM pi_scan_pages sp WHERE sp.post_id = s.post_id " "AND sp.exif_type IN (SELECT CAST(value AS INTEGER) " "FROM json_each(p.deep_scan_types)))) THEN 1 END) found " "FROM pi_searches p JOIN json_each(p.post_ids) j " "LEFT JOIN pi_scans s ON s.post_id = CAST(j.value AS TEXT) " f"WHERE p.id IN ({marks}) GROUP BY p.id" ") UPDATE pi_searches SET " "found_exif = COALESCE((SELECT found FROM counts WHERE counts.id = pi_searches.id), 0), " "total_searched = COALESCE((SELECT scanned FROM counts WHERE counts.id = pi_searches.id), 0) " "WHERE id IN (SELECT id FROM counts)", "args": [{"type": "text", "value": sid} for sid in batch], } ] ) updated += len(batch) return updated async def get_unscanned_post_ids(search_id): resp = await turso_execute( [ { "sql": "WITH ids AS (" "SELECT CAST(j.value AS TEXT) post_id, CAST(j.key AS INTEGER) pos " "FROM pi_searches p, json_each(p.post_ids) j WHERE p.id = ?" ") SELECT ids.post_id FROM ids " "LEFT JOIN pi_scans s ON s.post_id = ids.post_id " "WHERE s.post_id IS NULL ORDER BY ids.pos", "args": [{"type": "text", "value": search_id}], } ] ) return [cell_value(row, 0, "") for row in first_rows(resp) if cell_value(row, 0, "")] async def get_deep_scan_targets(search_id, targeted=False): start_page = "0" if targeted else "COALESCE(MAX(p.page), -1) + 1" condition = ( "s.exif_type IS NOT NULL" if targeted else "COALESCE(s.deep_scanned, 0) = 0 AND s.exif_type IS NOT NULL" ) resp = await turso_execute( [ { "sql": "WITH ids AS (" "SELECT CAST(j.value AS TEXT) post_id, CAST(j.key AS INTEGER) pos " "FROM pi_searches p, json_each(p.post_ids) j WHERE p.id = ?" f") SELECT ids.post_id, {start_page} start_page FROM ids " "JOIN pi_scans s ON s.post_id = ids.post_id " "LEFT JOIN pi_scan_pages p ON p.post_id = ids.post_id " f"WHERE {condition} " "GROUP BY ids.post_id, ids.pos ORDER BY ids.pos", "args": [{"type": "text", "value": search_id}], } ] ) return [ { "post_id": cell_value(row, 0, ""), "start_page": int(cell_value(row, 1, "0")), } for row in first_rows(resp) if cell_value(row, 0, "") ] async def get_incomplete_search_ids(): resp = await turso_execute( [ { "sql": "WITH counts AS (" "SELECT p.id, COUNT(j.value) total, COUNT(s.post_id) scanned " "FROM pi_searches p JOIN json_each(p.post_ids) j " "LEFT JOIN pi_scans s ON s.post_id = CAST(j.value AS TEXT) " "GROUP BY p.id" ") SELECT p.id FROM pi_searches p JOIN counts ON counts.id = p.id " "WHERE total > scanned OR (COALESCE(p.deep_scan, 0) = 1 AND EXISTS (" "SELECT 1 FROM json_each(p.post_ids) j JOIN pi_scans s " "ON s.post_id = CAST(j.value AS TEXT) " "WHERE COALESCE(s.deep_scanned, 0) = 0 " "AND s.exif_type IS NOT NULL)) ORDER BY p.id", } ] ) return [cell_value(row, 0, "") for row in first_rows(resp) if cell_value(row, 0, "")] async def search_uses_deep_scan(search_id): resp = await turso_execute( [ { "sql": "SELECT deep_scan FROM pi_searches WHERE id = ?", "args": [{"type": "text", "value": search_id}], } ] ) rows = first_rows(resp) return bool(rows and int(cell_value(rows[0], 0, "0"))) async def get_search_deep_scan_types(search_id): resp = await turso_execute( [ { "sql": "SELECT deep_scan_types FROM pi_searches WHERE id = ?", "args": [{"type": "text", "value": search_id}], } ] ) rows = first_rows(resp) return exif_codes(json.loads(cell_value(rows[0], 0, "[]"))) if rows else [] async def get_refresh_search_ids(): resp = await turso_execute( [ { "sql": "SELECT id FROM pi_searches " "WHERE post_ids != '[]' AND total_searched = 0 " "ORDER BY id DESC LIMIT ?", "args": [{"type": "integer", "value": str(REFRESH_SEARCH_LIMIT)}], } ] ) return [cell_value(row, 0, "") for row in first_rows(resp) if cell_value(row, 0, "")] async def get_result_counts(search_id): id_arg = {"type": "text", "value": search_id} resp = await turso_execute( [ { "sql": RESULT_BASE_SQL + "SELECT COUNT(*), COUNT(s.post_id), " "(SELECT COUNT(DISTINCT post_id) FROM tags) " "FROM ids LEFT JOIN pi_scans s ON s.post_id = ids.post_id", "args": [id_arg], }, { "sql": RESULT_BASE_SQL + "SELECT exif_type, COUNT(DISTINCT post_id) FROM tags GROUP BY exif_type " "UNION ALL SELECT 0, COUNT(*) FROM ids " "WHERE NOT EXISTS (SELECT 1 FROM tags WHERE tags.post_id = ids.post_id)", "args": [id_arg], }, ] ) count_rows = first_rows(resp) if not count_rows: return None type_rows = first_rows(resp, 1) return { "raw_total": int(cell_value(count_rows[0], 0, "0")), "scanned_count": int(cell_value(count_rows[0], 1, "0")), "found_exif": int(cell_value(count_rows[0], 2, "0")), "type_counts": { int(cell_value(row, 0, "0")): int(cell_value(row, 1, "0")) for row in type_rows }, } RESULT_BASE_SQL = ( "WITH ids AS (" "SELECT CAST(j.value AS TEXT) post_id, CAST(j.key AS INTEGER) + 1 scan_index, " "COALESCE(p.deep_scan_types, '[]') deep_scan_types " "FROM pi_searches p, json_each(p.post_ids) j WHERE p.id = ?" "), tags AS (" "SELECT p.post_id, p.exif_type FROM pi_scan_pages p JOIN ids ON ids.post_id = p.post_id " "WHERE p.exif_type IS NOT NULL AND (ids.deep_scan_types = '[]' OR EXISTS (" "SELECT 1 FROM json_each(ids.deep_scan_types) WHERE CAST(value AS INTEGER) = p.exif_type)) UNION " "SELECT s.post_id, s.exif_type FROM pi_scans s JOIN ids ON ids.post_id = s.post_id " "WHERE s.exif_type IS NOT NULL AND (ids.deep_scan_types = '[]' OR EXISTS (" "SELECT 1 FROM json_each(ids.deep_scan_types) WHERE CAST(value AS INTEGER) = s.exif_type))) " ) def result_filter_sql(allowed): if allowed is None: return "EXISTS (SELECT 1 FROM tags t WHERE t.post_id = s.post_id)", [] parts = [] args = [] if 0 in allowed: parts.append("NOT EXISTS (SELECT 1 FROM tags t WHERE t.post_id = s.post_id)") codes = [x for x in allowed if x] if codes: parts.append( "EXISTS (SELECT 1 FROM tags t WHERE t.post_id = s.post_id " f"AND t.exif_type IN ({','.join('?' for _ in codes)}))" ) args = [{"type": "integer", "value": str(x)} for x in codes] if not parts: return "0", [] return "(" + " OR ".join(parts) + ")", args def scan_filter_sql(allowed): parts = [] args = [] if 0 in allowed: parts.append("exif_type IS NULL") codes = [x for x in allowed if x] if codes: parts.append(f"exif_type IN ({','.join('?' for _ in codes)})") args = [{"type": "integer", "value": str(x)} for x in codes] if not parts: return "0", [] return "(" + " OR ".join(parts) + ")", args def exif_label(code): return EXIF_CODE_TO_TYPE.get(code) if code else "None" def encrypt_log(log): key = hashlib.sha256(RECLASS_LOG_KEY.encode()).digest() nonce = secrets.token_bytes(12) data = json.dumps(log, separators=(",", ":")).encode() encrypted = AESGCM(key).encrypt(nonce, data, None) return base64.b64encode(nonce + encrypted).decode() def decrypt_log(value): if not value: return [] raw = base64.b64decode(value) key = hashlib.sha256(RECLASS_LOG_KEY.encode()).digest() data = AESGCM(key).decrypt(raw[:12], raw[12:], None) return json.loads(data.decode()) async def get_result_pages(post_ids): if not post_ids: return {} pages = {} for i in range(0, len(post_ids), 500): chunk = post_ids[i : i + 500] marks = ",".join("?" for _ in chunk) resp = await turso_execute( [ { "sql": f"SELECT post_id, page, url, exif_type FROM pi_scan_pages WHERE post_id IN ({marks}) ORDER BY post_id, page", "args": [{"type": "text", "value": str(pid)} for pid in chunk], } ] ) for row in first_rows(resp): exif_type = cell_value(row, 3) pages.setdefault(cell_value(row, 0, ""), []).append( { "page": int(cell_value(row, 1, "0")), "url": cell_value(row, 2, ""), "exif_type": int(exif_type) if exif_type else None, } ) return pages async def get_download_targets(search_id, through, allowed): where_sql, filter_args = result_filter_sql(allowed) resp = await turso_execute( [ { "sql": RESULT_BASE_SQL + "SELECT ids.post_id, ids.scan_index, COALESCE(p.url, s.url) " "FROM ids JOIN pi_scans s ON s.post_id = ids.post_id " "LEFT JOIN pi_scan_pages p ON p.post_id = ids.post_id AND (" "ids.deep_scan_types = '[]' OR EXISTS (SELECT 1 FROM json_each(ids.deep_scan_types) " "WHERE CAST(value AS INTEGER) = p.exif_type)) " f"WHERE ids.scan_index <= ? AND {where_sql} " "ORDER BY ids.scan_index, COALESCE(p.page, 0)", "args": [ {"type": "text", "value": search_id}, {"type": "integer", "value": str(through)}, *filter_args, ], } ] ) return [ { "post_id": cell_value(row, 0, ""), "url": cell_value(row, 2, ""), } for row in first_rows(resp) if cell_value(row, 2, "") ] async def create_download_zip(targets): counts = {} for item in targets: pid = item["post_id"] counts[pid] = counts.get(pid, 0) + 1 fd, path = tempfile.mkstemp(suffix=".zip") os.close(fd) try: with zipfile.ZipFile(path, "w", compression=zipfile.ZIP_STORED) as archive: seen = {} for start in range(0, len(targets), DOWNLOAD_BATCH_SIZE): batch = targets[start : start + DOWNLOAD_BATCH_SIZE] rows = await asyncio.gather( *( fetch_cached_image( item["post_id"], page_num_from_url(item["url"]), item["url"] if item["url"].startswith("http") else IMG_BASE + item["url"].lstrip("/"), ) for item in batch ) ) for item, (data, _) in zip(batch, rows): pid = item["post_id"] seen[pid] = seen.get(pid, 0) + 1 ext = image_ext(item["url"]) if ext == "jpeg": ext = "jpg" if ext not in ("png", "jpg", "gif", "webp"): ext = "bin" suffix = f"_{seen[pid]}" if counts[pid] > 1 else "" archive.writestr(f"{pid}{suffix}.{ext}", data) return path except Exception: Path(path).unlink(missing_ok=True) raise async def get_result_page(search_id, page, allowed, sort_bookmarks=True): where_sql, filter_args = result_filter_sql(allowed) page = max(page, 1) offset = (page - 1) * PAGE_SIZE id_arg = {"type": "text", "value": search_id} order_sql = ( "s.bookmark_count IS NULL, s.bookmark_count DESC, ids.scan_index" if sort_bookmarks else "ids.scan_index" ) resp = await turso_execute( [ { "sql": RESULT_BASE_SQL + "SELECT COUNT(*), COUNT(s.post_id), " "(SELECT COUNT(DISTINCT post_id) FROM tags), MAX(ids.deep_scan_types) " "FROM ids LEFT JOIN pi_scans s ON s.post_id = ids.post_id", "args": [id_arg], }, { "sql": RESULT_BASE_SQL + "SELECT COUNT(*) FROM ids JOIN pi_scans s ON s.post_id = ids.post_id " f"WHERE {where_sql}", "args": [id_arg, *filter_args], }, { "sql": RESULT_BASE_SQL + "SELECT ids.post_id, ids.scan_index, s.url, s.exif_type, " "s.bookmark_count, s.view_count " "FROM ids JOIN pi_scans s ON s.post_id = ids.post_id " f"WHERE {where_sql} ORDER BY {order_sql} LIMIT ? OFFSET ?", "args": [ id_arg, *filter_args, {"type": "integer", "value": str(PAGE_SIZE)}, {"type": "integer", "value": str(offset)}, ], }, { "sql": RESULT_BASE_SQL + "SELECT exif_type, COUNT(DISTINCT post_id) FROM tags GROUP BY exif_type " "UNION ALL SELECT 0, COUNT(*) FROM ids " "WHERE NOT EXISTS (SELECT 1 FROM tags WHERE tags.post_id = ids.post_id)", "args": [id_arg], }, ] ) count_rows = first_rows(resp) if not count_rows: return None total_rows = first_rows(resp, 1) item_rows = first_rows(resp, 2) type_rows = first_rows(resp, 3) raw_total = int(cell_value(count_rows[0], 0, "0")) scanned_count = int(cell_value(count_rows[0], 1, "0")) total = int(cell_value(total_rows[0], 0, "0")) if total_rows else 0 type_counts = { int(cell_value(row, 0, "0")): int(cell_value(row, 1, "0")) for row in type_rows } deep_scan_types = set(json.loads(cell_value(count_rows[0], 3, "[]"))) result_pages = await get_result_pages([cell_value(row, 0, "") for row in item_rows]) items = [] for row in item_rows: pid = cell_value(row, 0, "") url = cell_value(row, 2, "") exif_type = cell_value(row, 3) bookmark_count = cell_value(row, 4) view_count = cell_value(row, 5) page_rows = [ item for item in result_pages.get(pid) or [] if not deep_scan_types or item["exif_type"] in deep_scan_types ] exif_types = sorted( {item["exif_type"] for item in page_rows if item["exif_type"]} ) if not exif_types and exif_type and ( not deep_scan_types or int(exif_type) in deep_scan_types ): exif_types = [int(exif_type)] enabled = set(allowed if allowed is not None else exif_types) visible_rows = [ item for item in page_rows if (item["exif_type"] or 0) in enabled and (item["exif_type"] or not exif_types) ] selected = visible_rows[0] if visible_rows else None if selected: url = selected["url"] pages = [item["page"] for item in visible_rows] or [page_num_from_url(url)] page_types = { str(item["page"]): item["exif_type"] or 0 for item in page_rows } if not page_rows: page_types[str(pages[0])] = int(exif_type) if exif_type else 0 items.append( { "post_id": pid, "scan_index": int(cell_value(row, 1, "0")), "url": url, "exif_type": int(exif_type) if exif_type else None, "bookmark_count": int(bookmark_count) if bookmark_count is not None else None, "view_count": int(view_count) if view_count is not None else None, "exif_types": exif_types, "scanned": True, "pages": pages, "page_types": page_types, "page_count": len(pages), **image_links(pid, url), } ) return { "items": items, "total": total, "page": page, "page_size": PAGE_SIZE, "pages": max((total + PAGE_SIZE - 1) // PAGE_SIZE, 1), "raw_total": raw_total, "scanned_count": scanned_count, "type_counts": type_counts, } def user_search_id(user_id): return f"{base26_time()}_{user_id}" async def pixiv_search_live( url, pages, mode, phpsessid, search_id, deep_scan=True, deep_scan_types=None ): keywords = get_search_keywords(url) api_url = get_search_api_url(url, keywords) first_url = f"{api_url}&p=1" cookies = {"PHPSESSID": phpsessid} post_ids = [] seen = set() done = 0 async def get_page(page): return page, await fetch_page(session, f"{api_url}&p={page}") async with aiohttp.ClientSession(cookies=cookies, headers=PIXIV_HEADERS) as session: tasks = [ asyncio.create_task(get_page(p)) for p in range(1, pages + 1) ] page_data = {} next_page = 1 try: for coro in asyncio.as_completed(tasks): if task_canceled(search_id): raise TaskCanceled() page, data = await coro done += 1 if search_id in ACTIVE_TASKS: ACTIVE_TASKS[search_id].update({"total": pages, "done": done}) page_data[page] = data while next_page in page_data: data = page_data.pop(next_page) next_page += 1 if data.get("error"): continue body = data.get("body") or {} posts = (body.get("illustManga") or {}).get("data") or [] if mode == "ai": posts = [p for p in posts if is_ai_post(p)] elif mode == "real": posts = [p for p in posts if not is_ai_post(p)] for post in posts: pid = str(post.get("id") or "") if pid and pid not in seen: seen.add(pid) post_ids.append(pid) if done % 25 == 0: if task_canceled(search_id): raise TaskCanceled() await save_search( search_id, post_ids, deep_scan, deep_scan_types ) finally: for task in tasks: if not task.done(): task.cancel() await asyncio.gather(*tasks, return_exceptions=True) if task_canceled(search_id): raise TaskCanceled() await save_search(search_id, post_ids, deep_scan, deep_scan_types) return post_ids, keywords, first_url async def pixiv_user_posts(user_id, phpsessid): cookies = {"PHPSESSID": phpsessid} async with aiohttp.ClientSession(cookies=cookies, headers=PIXIV_HEADERS) as session: data = await fetch_page( session, f"https://www.pixiv.net/ajax/user/{user_id}/profile/all" ) body = data.get("body") or {} posts = list((body.get("illusts") or {}).keys()) return {"user_id": user_id, "post_ids": posts} async def fetch_page(session, url): async with session.get(url) as r: return await r.json() async def get_post_pages(post_id, session): key = str(post_id) cached = PAGE_URL_CACHE.get(key) now = time.time() if cached and now - cached["time"] < PAGE_URL_CACHE_MAX_AGE: return cached["pages"] inflight = PAGE_URL_INFLIGHT.get(key) if inflight: return await inflight async def load(): data = await fetch_page( session, f"https://www.pixiv.net/ajax/illust/{post_id}/pages" ) pages = [] for page in data.get("body") or []: urls = page.get("urls") or {} pages.append( { "original": urls.get("original") or "", "regular": urls.get("regular") or "", "small": urls.get("small") or "", "thumb_mini": urls.get("thumb_mini") or "", } ) PAGE_URL_CACHE[key] = {"time": time.time(), "pages": pages} if len(PAGE_URL_CACHE) > PAGE_URL_CACHE_LIMIT: cutoff = time.time() - PAGE_URL_CACHE_MAX_AGE for cid, row in list(PAGE_URL_CACHE.items()): if row["time"] < cutoff: PAGE_URL_CACHE.pop(cid, None) return pages task = asyncio.create_task(load()) PAGE_URL_INFLIGHT[key] = task try: return await task finally: if PAGE_URL_INFLIGHT.get(key) is task: PAGE_URL_INFLIGHT.pop(key, None) async def fetch_post_popularity(post_id, session): try: data = await fetch_page( session, f"https://www.pixiv.net/ajax/illust/{post_id}" ) body = data.get("body") or {} return post_id, int(body["bookmarkCount"]), int(body["viewCount"]) except Exception: return None async def get_missing_popularity_post_ids(search_id): resp = await turso_execute( [ { "sql": "WITH ids AS (" "SELECT CAST(j.value AS TEXT) post_id, CAST(j.key AS INTEGER) pos " "FROM pi_searches p, json_each(p.post_ids) j WHERE p.id = ?" ") SELECT ids.post_id FROM ids " "JOIN pi_scans s ON s.post_id = ids.post_id " "WHERE s.bookmark_count IS NULL OR s.view_count IS NULL " "ORDER BY ids.pos", "args": [{"type": "text", "value": search_id}], } ] ) return [cell_value(row, 0, "") for row in first_rows(resp)] async def cache_search_popularity(search_id): post_ids = await get_missing_popularity_post_ids(search_id) if not post_ids: return session = await pixiv_client() for i in range(0, len(post_ids), POPULARITY_BATCH_SIZE): chunk = post_ids[i : i + POPULARITY_BATCH_SIZE] results = await asyncio.gather( *(fetch_post_popularity(post_id, session) for post_id in chunk) ) stmts = [ { "sql": "UPDATE pi_scans SET " "bookmark_count = COALESCE(bookmark_count, ?), " "view_count = COALESCE(view_count, ?) WHERE post_id = ?", "args": [ {"type": "integer", "value": str(result[1])}, {"type": "integer", "value": str(result[2])}, {"type": "text", "value": str(result[0])}, ], } for result in results if result ] if stmts: await turso_batch(stmts) def decode_meta_text(data): for enc in ("utf-8", "latin-1"): try: return data.decode(enc).strip() except UnicodeDecodeError: pass return data.decode("utf-8", errors="replace").strip() def parse_png_metadata(data): if not data.startswith(b"\x89PNG\r\n\x1a\n"): return None fields = {} index = 8 while index < len(data): if index + 8 > len(data): break chunk_len = int.from_bytes(data[index : index + 4], "big") chunk_type = data[index + 4 : index + 8].decode("ascii", errors="ignore") index += 8 content = data[index : min(index + chunk_len, len(data))] if chunk_type == "tEXt" and b"\0" in content: key, text = content.split(b"\0", 1) fields[decode_meta_text(key)] = decode_meta_text(text) elif chunk_type == "iTXt" and b"\0" in content: parts = content.split(b"\0", 5) if len(parts) == 6: key, flag, method, _, _, text = parts if flag == b"\x01" and method == b"\x00": try: text = zlib.decompress(text) except zlib.error: text = b"" fields[decode_meta_text(key)] = decode_meta_text(text) elif chunk_type == "zTXt" and b"\0" in content: key, rest = content.split(b"\0", 1) if rest[:1] == b"\x00": try: fields[decode_meta_text(key)] = decode_meta_text( zlib.decompress(rest[1:]) ) except zlib.error: pass index += chunk_len + 4 return fields or None def decode_exif_text(value): if value is None: return "" if isinstance(value, bytes): prefixes = ( (b"ASCII\x00\x00\x00", "utf-8"), (b"UNICODE\x00", "utf-16-be"), (b"JIS\x00\x00\x00\x00\x00", "shift_jis"), ) for prefix, enc in prefixes: if value.startswith(prefix): try: return value[len(prefix) :].rstrip(b"\0").decode(enc).strip() except UnicodeDecodeError: break return decode_meta_text(value.rstrip(b"\0")) return str(value).strip() def add_exif_fields(fields, exif): for tag, value in exif.items(): name = ExifTags.TAGS.get(tag, str(tag)) fields[EXIF_FIELD_NAMES.get(name, name)] = decode_exif_text(value) def parse_pil_metadata(data): fields = {} try: image = Image.open(io.BytesIO(data)) for key, value in (image.info or {}).items(): if isinstance(value, (str, bytes)): fields[str(key)] = decode_exif_text(value) exif = image.getexif() add_exif_fields(fields, exif) if hasattr(ExifTags, "IFD"): add_exif_fields(fields, exif.get_ifd(ExifTags.IFD.Exif)) except Exception: pass return {k: v for k, v in fields.items() if v} or None def parse_jpeg_comments(data): if not data.startswith(b"\xff\xd8"): return None fields = {} index = 2 n = 1 while index + 4 <= len(data): if data[index] != 0xFF: index += 1 continue marker = data[index + 1] index += 2 if marker == 0xDA: break if marker in (0x01, *range(0xD0, 0xD9)): continue size = int.from_bytes(data[index : index + 2], "big") if size < 2 or index + size > len(data): break content = data[index + 2 : index + size] if marker == 0xFE: fields[f"Comment{n}"] = decode_meta_text(content) n += 1 index += size return fields or None def parse_xmp_metadata(data): start = data.find(b"", start) if end >= 0: end += len(b"") else: end = data.find(b"?>", start) if end >= 0: end += 2 if end < 0: return None return {"XMP": decode_meta_text(data[start:end])} def parse_image_metadata(data): fields = {} if data.startswith(b"\x89PNG\r\n\x1a\n"): parsers = (parse_png_metadata, parse_xmp_metadata) elif data.startswith(b"\xff\xd8"): parsers = (parse_pil_metadata, parse_jpeg_comments, parse_xmp_metadata) elif data.startswith(b"RIFF") and data[8:12] == b"WEBP": parsers = (parse_pil_metadata, parse_xmp_metadata) else: parsers = (parse_pil_metadata, parse_xmp_metadata) for parser in parsers: parsed = parser(data) if parsed: fields.update(parsed) if not fields and data.startswith(b"\x89PNG\r\n\x1a\n"): parsed = parse_pil_metadata(data) if parsed: fields.update(parsed) return fields or None def parse_json_maybe(value): try: return json.loads(value) except (TypeError, json.JSONDecodeError): return None def has_novelai_comment(text): data = text if isinstance(text, dict) else parse_json_maybe(text) if isinstance(data, dict): keys = set(data) nai_keys = { "cfg_rescale", "legacy_v3_extend", "noise_schedule", "request_type", "signed_hash", "uc", "v4_negative_prompt", "v4_prompt", } if keys & nai_keys: return True return ( isinstance(data.get("prompt"), str) and isinstance(data.get("steps"), int) and ("scale" in data or "sampler" in data) ) if not isinstance(text, str): return False text = text[:4096] needles = ('"prompt"', '"steps"', '"sampler"', '"scale"', '"uc"', '"request_type"') return all(needle in text for needle in needles[:4]) or ( '"prompt"' in text and ('"signed_hash"' in text or '"v4_prompt"' in text) ) def has_novelai_vibe(fields): for value in (fields, *fields.values()): data = value if isinstance(value, dict) else parse_json_maybe(value) if isinstance(data, dict): refs = data.get("reference_image_multiple") if isinstance(refs, list) and refs: return True return False def has_comfy_metadata(fields): for key in ("prompt", "workflow"): data = parse_json_maybe(fields.get(key)) if isinstance(data, dict): values = data.values() if key == "prompt" else [data] for item in values: if isinstance(item, dict) and ( "class_type" in item or "last_node_id" in item ): return True return False def field_json_values(fields): for value in fields.values(): data = parse_json_maybe(value) if isinstance(data, dict): yield data def has_fooocus_metadata(fields): scheme = fields.get("fooocus_scheme") or fields.get("MakerNote") if str(scheme).casefold() in ("fooocus", "a1111"): return True software = fields.get("Software", "").casefold() if "fooocus" in software: return True for data in field_json_values(fields): keys = set(data) if { "adaptive_cfg", "base_model_hash", "full_negative_prompt", "full_prompt", "loras", "refiner_model_hash", } & keys: return True return False def has_invoke_metadata(fields): if any(k.casefold() in ("dream", "sd-metadata") for k in fields): return True joined = "\n".join(str(v) for v in fields.values()).casefold() if "invokeai" in joined or "invoke ai" in joined: return True for data in field_json_values(fields): app = str(data.get("app_id", "")).casefold() if "invoke" in app: return True return False def has_swarm_metadata(fields): for data in field_json_values(fields): if {"sui_image_params", "sui_extra_data", "sui_models"} & set(data): return True return "sui_image_params" in "\n".join(str(v) for v in fields.values()).casefold() def has_sd_metadata(fields): if "parameters" in fields: return True for value in fields.values(): text = str(value) folded = text.casefold() if ( "steps:" in folded and "seed:" in folded and ("sampler:" in folded or "cfg scale:" in folded or "model:" in folded) ): return True data = parse_json_maybe(value) if isinstance(data, dict): keys = set(data) if "prompt" in keys and "seed" in keys and ( "steps" in keys or "sampler" in keys or "cfg_scale" in keys ): return True return False def determine_exif_type(fields): if not fields: return None values = [str(v) for v in fields.values()] joined = "\n".join(values).casefold() software = fields.get("Software", "").casefold() source = fields.get("Source", "").casefold() title = fields.get("Title", "") comment = fields.get("Comment", "") if ( "novelai" in software or "novelai" in source or title == "AI generated image" or has_novelai_comment(comment) or has_novelai_comment(fields) ): return "nai_vibe" if has_novelai_vibe(fields) else "nai" if has_comfy_metadata(fields): return "comfy" if has_fooocus_metadata(fields): return "fooocus" if has_invoke_metadata(fields): return "invoke" if has_swarm_metadata(fields): return "swarm" if has_sd_metadata(fields) or joined.startswith("parameters"): return "sd" if "midjourney" in joined: return "mj" if "celsys" in software: return "celsys" if "photoshop" in software or "adobe" in joined: return "photoshop" return None def image_ext(url): return urlsplit(url).path.rsplit(".", 1)[-1].casefold() def byteize(bits): bits = bits[: (bits.shape[0] // 8) * 8] bits = np.bitwise_and(bits, 1) bits = bits.reshape((-1, 8)) return np.packbits(bits, axis=1).flatten() def parse_stealth_channel(bits, sigs): raw = bytes(byteize(bits)) magic = next((sig for sig in sigs if raw.startswith(sig)), None) if not magic: return None, False try: start = len(magic) bit_len = int.from_bytes(raw[start : start + 4], "big") start += 4 if bit_len <= 0 or bit_len % 8 or start + bit_len // 8 > len(raw): return None, True payload = raw[start : start + bit_len // 8] if magic.endswith(b"comp"): payload = zlib.decompress(payload, wbits=31) text = decode_meta_text(payload) fields = parse_json_maybe(text) return (fields if isinstance(fields, dict) else {"Stealth": text}), True except (UnicodeDecodeError, ValueError, zlib.error): return None, True def parse_stealth_metadata(data): try: image = Image.open(io.BytesIO(data)) detected = False if "A" in image.getbands(): alpha = np.array(image.getchannel("A")).T.reshape((-1,)) fields, found = parse_stealth_channel(alpha, STEALTH_ALPHA_SIGS) if fields: return fields, True detected = found rgb = np.array(image.convert("RGB")).transpose((1, 0, 2)).reshape((-1,)) fields, found = parse_stealth_channel(rgb, STEALTH_RGB_SIGS) return fields, detected or found except Exception: return None, False async def scan_post(post_id, session, post_sem, img_sem): async with post_sem: try: pages = await get_post_pages(post_id, session) if not pages: return None image_urls = [(page, p["original"]) for page, p in enumerate(pages)] for page, url in image_urls: img_data = await fetch_image(session, url, img_sem, post_id, page) if not img_data: return None code = await asyncio.to_thread(scan_image_data, img_data) if code: return post_id, url, code, page == len(pages) - 1 return post_id, None, None, True except Exception: return None def scan_image_data(data, allowed=None): allowed = set(allowed) if allowed else None fallback = None exif_type = determine_exif_type(parse_image_metadata(data)) code = EXIF_TYPE_TO_CODE.get(exif_type) if exif_type in FALLBACK_EXIF_TYPES: if allowed is None or code in allowed: fallback = code elif code and (allowed is None or code in allowed): return code metadata, detected = parse_stealth_metadata(data) exif_type = determine_exif_type(metadata) code = EXIF_TYPE_TO_CODE.get(exif_type) if exif_type in FALLBACK_EXIF_TYPES: if allowed is None or code in allowed: fallback = fallback or code elif exif_type == "nai": code = EXIF_TYPE_TO_CODE.get("nai_stealth") if allowed is None or code in allowed: return code elif code and (allowed is None or code in allowed): return code code = EXIF_TYPE_TO_CODE.get("nai_stealth") return code if detected and (allowed is None or code in allowed) else fallback async def scan_page_result(post_id, url, session, img_sem, allowed=None): page = page_num_from_url(url) img_data = await fetch_image(session, url, img_sem, post_id, page) if not img_data: return None exif_type = await asyncio.to_thread(scan_image_data, img_data, allowed) return post_id, url, exif_type async def deep_scan_post(item, session, post_sem, img_sem, allowed=None): async with post_sem: try: post_id = item["post_id"] start_page = item["start_page"] pages = await get_post_pages(post_id, session) if not pages: return None results = [] fresh = [] complete = True for page, urls in enumerate(pages[start_page:], start_page): url = urls["original"] try: result = await scan_page_result( post_id, url, session, img_sem, allowed ) except Exception: result = None if not result: complete = False break row = (post_id, page, result[1], result[2]) results.append(row) fresh.append(row) return post_id, results, fresh, complete except Exception: return None async def fetch_image(session, url, sem, post_id, page): if sem: async with sem: return (await fetch_cached_image(post_id, page, url, session))[0] return (await fetch_cached_image(post_id, page, url, session))[0] async def run_scan(post_ids, phpsessid, task_id=None, save_live=False, done_offset=0): cookies = {"PHPSESSID": phpsessid} results = [] completed = 0 async with aiohttp.ClientSession(cookies=cookies, headers=PIXIV_HEADERS) as session: post_iter = iter(post_ids) tasks = set() created = set() def fill_tasks(): while len(tasks) < POST_SCAN_LIMIT: if task_canceled(task_id): raise TaskCanceled() try: pid = next(post_iter) except StopIteration: break task = asyncio.create_task( scan_post(pid, session, POST_SCAN_SEM, FULL_IMAGE_SEM) ) tasks.add(task) created.add(task) if task_id: TASK_SCAN_TASKS.setdefault(task_id, set()).add(task) pending = [] try: fill_tasks() while tasks: if task_canceled(task_id): raise TaskCanceled() done, tasks = await asyncio.wait( tasks, timeout=1, return_when=asyncio.FIRST_COMPLETED ) if not done: continue if task_canceled(task_id): raise TaskCanceled() for task in done: if task.cancelled(): raise TaskCanceled() result = task.result() created.discard(task) if task_id: TASK_SCAN_TASKS.get(task_id, set()).discard(task) completed += 1 if result: results.append(result) pending.append(result) if task_id and task_id in ACTIVE_TASKS: ACTIVE_TASKS[task_id]["done"] = done_offset + completed if save_live and len(pending) >= 20: if task_canceled(task_id): raise TaskCanceled() await save_scan_results(pending) pending = [] fill_tasks() if save_live and pending: if task_canceled(task_id): raise TaskCanceled() await save_scan_results(pending) finally: for task in created: if not task.done(): task.cancel() await asyncio.gather(*created, return_exceptions=True) if task_id and task_id in TASK_SCAN_TASKS: TASK_SCAN_TASKS[task_id].difference_update(created) if not TASK_SCAN_TASKS[task_id]: TASK_SCAN_TASKS.pop(task_id, None) return results async def run_deep_scan(targets, phpsessid, task_id=None, allowed=None): cookies = {"PHPSESSID": phpsessid} completed = 0 pending = [] async with aiohttp.ClientSession(cookies=cookies, headers=PIXIV_HEADERS) as session: target_iter = iter(targets) tasks = set() created = set() def fill_tasks(): while len(tasks) < POST_SCAN_LIMIT: if task_canceled(task_id): raise TaskCanceled() try: item = next(target_iter) except StopIteration: break task = asyncio.create_task( deep_scan_post( item, session, POST_SCAN_SEM, FULL_IMAGE_SEM, allowed ) ) tasks.add(task) created.add(task) if task_id: TASK_SCAN_TASKS.setdefault(task_id, set()).add(task) try: fill_tasks() while tasks: if task_canceled(task_id): raise TaskCanceled() done, tasks = await asyncio.wait( tasks, timeout=1, return_when=asyncio.FIRST_COMPLETED ) if task_canceled(task_id): raise TaskCanceled() for task in done: result = task.result() created.discard(task) if task_id: TASK_SCAN_TASKS.get(task_id, set()).discard(task) completed += 1 if result: pending.append(result) if task_id and task_id in ACTIVE_TASKS: ACTIVE_TASKS[task_id]["done"] = completed if len(pending) >= 20: await save_deep_scan_results(pending, bool(allowed)) pending = [] fill_tasks() if pending: await save_deep_scan_results(pending, bool(allowed)) finally: for task in created: if not task.done(): task.cancel() await asyncio.gather(*created, return_exceptions=True) if task_id and task_id in TASK_SCAN_TASKS: TASK_SCAN_TASKS[task_id].difference_update(created) if not TASK_SCAN_TASKS[task_id]: TASK_SCAN_TASKS.pop(task_id, None) async def save_scan_results(results): stmts = [] for post_id, url, exif_type, deep_scanned in results: short_url = url.replace(IMG_BASE, "", 1) if url else "" stmts.append( { "sql": "DELETE FROM pi_scan_pages WHERE post_id = ?", "args": [{"type": "text", "value": str(post_id)}], } ) stmts.append( { "sql": "INSERT INTO pi_scans " "(post_id, url, exif_type, deep_scanned) VALUES (?, ?, ?, ?) " "ON CONFLICT(post_id) DO UPDATE SET " "url = excluded.url, exif_type = excluded.exif_type, " "deep_scanned = excluded.deep_scanned", "args": [ {"type": "text", "value": str(post_id)}, {"type": "text", "value": short_url}, {"type": "integer", "value": str(exif_type)} if exif_type else {"type": "null"}, {"type": "integer", "value": "1" if deep_scanned else "0"}, ], } ) if exif_type: stmts.append( { "sql": "INSERT OR REPLACE INTO pi_scan_pages " "(post_id, page, url, exif_type) VALUES (?, ?, ?, ?)", "args": [ {"type": "text", "value": str(post_id)}, {"type": "integer", "value": str(page_num_from_url(url))}, {"type": "text", "value": short_url}, {"type": "integer", "value": str(exif_type)}, ], } ) if stmts: for i in range(0, len(stmts), 200): await turso_batch(stmts[i : i + 200]) async def save_deep_scan_results(results, preserve_empty=False): stmts = [] stored = await get_result_pages([str(result[0]) for result in results]) for post_id, pages, fresh, complete in results: for _, page, url, exif_type in fresh: if preserve_empty and not exif_type: continue short_url = url.replace(IMG_BASE, "", 1) if url else "" stmts.append( { "sql": "INSERT OR REPLACE INTO pi_scan_pages (post_id, page, url, exif_type) VALUES (?, ?, ?, ?)", "args": [ {"type": "text", "value": str(post_id)}, {"type": "integer", "value": str(page)}, {"type": "text", "value": short_url}, {"type": "integer", "value": str(exif_type)} if exif_type else {"type": "null"}, ], } ) page_map = { row["page"]: (post_id, row["page"], row["url"], row["exif_type"]) for row in stored.get(str(post_id), ()) } page_map.update( { row[1]: row for row in pages if row[3] or not preserve_empty } ) all_pages = list(page_map.values()) counts = {} for _, _, _, exif_type in all_pages: if not exif_type: continue counts[exif_type] = counts.get(exif_type, 0) + 1 best = max(counts, key=lambda code: (counts[code], -code)) if counts else None url = next((row[2] for row in all_pages if row[3] == best), "") short_url = url.replace(IMG_BASE, "", 1) if url else "" if not complete and not best: stmts.append( { "sql": "UPDATE pi_scans SET deep_scanned = 0 WHERE post_id = ?", "args": [{"type": "text", "value": str(post_id)}], } ) continue stmts.append( { "sql": "UPDATE pi_scans SET url = ?, exif_type = ?, deep_scanned = ? WHERE post_id = ?", "args": [ {"type": "text", "value": short_url}, {"type": "integer", "value": str(best)} if best else {"type": "null"}, {"type": "integer", "value": "1" if complete else "0"}, {"type": "text", "value": str(post_id)}, ], } ) for i in range(0, len(stmts), 200): await turso_batch(stmts[i : i + 200]) async def get_reclass_targets(allowed): where_sql, args = scan_filter_sql(allowed) resp = await turso_execute( [ { "sql": f"SELECT post_id, url, exif_type FROM pi_scans WHERE {where_sql} ORDER BY post_id", "args": args, } ] ) items = [] for row in first_rows(resp): exif_type = cell_value(row, 2) items.append( { "post_id": cell_value(row, 0, ""), "url": cell_value(row, 1, "") or "", "exif_type": int(exif_type) if exif_type else None, } ) return items async def save_reclass_job(job_id, post_ids, filters): summary = { "scanned": 0, "changed": 0, "updated": 0, "filters": [exif_label(code) for code in filters], } await turso_execute( [ { "sql": "INSERT OR REPLACE INTO pi_reclass_jobs " "(id, post_ids, filters, total, scanned, changed, summary, log_enc, created_at) " "VALUES (?, ?, ?, ?, 0, 0, ?, ?, ?)", "args": [ {"type": "text", "value": job_id}, {"type": "text", "value": json.dumps(post_ids)}, {"type": "text", "value": json.dumps(filters)}, {"type": "integer", "value": str(len(post_ids))}, {"type": "text", "value": json.dumps(summary)}, {"type": "text", "value": encrypt_log([])}, {"type": "integer", "value": str(int(time.time()))}, ], } ] ) await patch_explorer_cache(job_id, front=True) async def finish_reclass_job(job_id, scanned, changes, updated, filters): summary = { "scanned": scanned, "changed": len(changes), "updated": updated, "filters": [exif_label(code) for code in filters], } await turso_execute( [ { "sql": "UPDATE pi_reclass_jobs SET scanned = ?, changed = ?, summary = ?, log_enc = ? WHERE id = ?", "args": [ {"type": "integer", "value": str(scanned)}, {"type": "integer", "value": str(len(changes))}, {"type": "text", "value": json.dumps(summary)}, {"type": "text", "value": encrypt_log(changes)}, {"type": "text", "value": job_id}, ], } ] ) await patch_explorer_cache(job_id) async def update_searches_for_post_ids(post_ids): if not post_ids: return ids = set() for chunk in [post_ids[i : i + 250] for i in range(0, len(post_ids), 250)]: placeholders = ",".join("?" for _ in chunk) resp = await turso_execute( [ { "sql": "SELECT DISTINCT p.id FROM pi_searches p, json_each(p.post_ids) j " f"WHERE CAST(j.value AS TEXT) IN ({placeholders})", "args": [{"type": "text", "value": str(pid)} for pid in chunk], } ] ) ids.update(cell_value(row, 0, "") for row in first_rows(resp)) updated = False for search_id in ids: if search_id: await update_search_stats(search_id) updated = True if updated: await rebuild_explorer_cache() async def get_scanned_post_ids(post_ids): if not post_ids: return {} chunks = [post_ids[i : i + 500] for i in range(0, len(post_ids), 500)] scanned = {} stmts = [] for chunk in chunks: placeholders = ",".join("?" for _ in chunk) stmts.append( { "sql": f"SELECT post_id, url, exif_type, bookmark_count, view_count " f"FROM pi_scans WHERE post_id IN ({placeholders})", "args": [{"type": "text", "value": str(pid)} for pid in chunk], } ) resp = await turso_execute(stmts) for result in resp.get("results") or []: if "response" not in result: continue rows = result["response"].get("result", {}).get("rows", []) for row in rows: pid = row[0].get("value") url_val = row[1].get("value") if row[1].get("type") != "null" else "" et = row[2].get("value") if row[2].get("type") != "null" else None bookmarks = row[3].get("value") if row[3].get("type") != "null" else None views = row[4].get("value") if row[4].get("type") != "null" else None scanned[pid] = { "url": url_val, "exif_type": int(et) if et else None, "bookmark_count": int(bookmarks) if bookmarks is not None else None, "view_count": int(views) if views is not None else None, } return scanned def cleanup_thumbs(): global LAST_THUMB_CLEANUP now = time.time() if now - LAST_THUMB_CLEANUP < THUMB_CLEAN_INTERVAL: return LAST_THUMB_CLEANUP = now THUMB_DIR.mkdir(exist_ok=True) for path in THUMB_DIR.glob("*.webp"): try: if now - path.stat().st_atime > THUMB_MAX_AGE: path.unlink() except OSError: pass def page_num_from_url(url): if not url: return 0 name = url.rsplit("/", 1)[-1] if "_p" not in name: return 0 try: return int(name.rsplit("_p", 1)[1].split(".", 1)[0]) except ValueError: return 0 def media_type_from_url(url): ext = urlsplit(url).path.rsplit(".", 1)[-1].casefold() if ext in ("jpg", "jpeg"): return "image/jpeg" if ext == "png": return "image/png" if ext == "gif": return "image/gif" if ext == "webp": return "image/webp" return "application/octet-stream" async def get_pixiv_image_url(post_id, page, size, phpsessid): session = await pixiv_client() pages = await get_post_pages(post_id, session) if not pages: raise HTTPException(status_code=404, detail="image not found") page = min(max(page, 0), len(pages) - 1) urls = pages[page] if size in ("full", "orig"): url = urls.get("original") or urls.get("regular") or urls.get("small") elif size == "thumb": url = ( urls.get("small") or urls.get("thumb_mini") or urls.get("regular") or urls.get("original") ) else: url = urls.get("regular") or urls.get("small") or urls.get("original") if not url: raise HTTPException(status_code=404, detail="image not found") return url def image_cache_path(post_id, page): return IMAGE_CACHE_DIR / str(post_id) / f"{page}.epng" def cleanup_image_cache_root(): IMAGE_CACHE_DIR.mkdir(parents=True, exist_ok=True) for path in IMAGE_CACHE_DIR.iterdir(): if path.is_file(): path.unlink() def read_image_cache(path): raw = path.read_bytes() return AESGCM(IMAGE_CACHE_KEY).decrypt(raw[:12], raw[12:], None) def write_image_cache(path, data): path.parent.mkdir(parents=True, exist_ok=True) nonce = secrets.token_bytes(12) data = nonce + AESGCM(IMAGE_CACHE_KEY).encrypt(nonce, data, None) fd, tmp = tempfile.mkstemp(dir=path.parent, suffix=".tmp") try: with os.fdopen(fd, "wb") as out: out.write(data) os.replace(tmp, path) except Exception: Path(tmp).unlink(missing_ok=True) raise async def fetch_cached_image(post_id, page, url, session=None): path = image_cache_path(post_id, page) lock = IMAGE_CACHE_LOCKS[hash(path) % len(IMAGE_CACHE_LOCKS)] async with lock: if path.exists(): try: data = await asyncio.to_thread(read_image_cache, path) return data, media_type_from_url(url) except Exception: pass if not url: url = await get_pixiv_image_url(post_id, page, "orig", PHPSESSID) client = session or await pixiv_client() async with client.get(url) as response: if response.status >= 400: raise HTTPException( status_code=response.status, detail="image fetch failed" ) data = await response.read() content_type = response.headers.get("Content-Type") or media_type_from_url(url) await asyncio.to_thread(write_image_cache, path, data) return data, content_type def write_webp(data, out, kind): image = Image.open(io.BytesIO(data)) if kind == "t": image.thumbnail((THUMB_SIZE, THUMB_SIZE)) else: image = image.resize( (max(int(image.width * WEBP_SCALE), 1), max(int(image.height * WEBP_SCALE), 1)) ) if image.mode not in ("RGB", "RGBA"): image = image.convert("RGB") image.save(out, "WEBP", quality=WEBP_QUALITY) async def create_webp(post_id, image_url, page=0, kind="t"): await asyncio.to_thread(cleanup_thumbs) scale_tag = THUMB_SIZE if kind == "t" else int(WEBP_SCALE * 100) out = THUMB_DIR / f"{post_id}_p{page}_{kind}{scale_tag}.webp" if out.exists(): os.utime(out, None) return out data, _ = await fetch_cached_image(post_id, page, image_url) if not data: raise HTTPException(status_code=404, detail="image not found") await asyncio.to_thread(write_webp, data, out, kind) return out def image_links(post_id, url): page = page_num_from_url(url) suffix = f"?p={page}&v={IMAGE_URL_VERSION}" pid = quote(str(post_id), safe="") thumb_url = f"/api/i/{pid}/t{suffix}" preview_url = f"/api/i/{pid}/v{suffix}" return { "image_url": thumb_url, "preview_url": preview_url, "download_url": f"/api/i/{pid}/o{suffix}", "full_image_url": preview_url, "page": page, } async def bg_search_task(search_id, url, pages, mode, phpsessid): if not start_task(search_id, { "type": "search", "phase": "searching", "total": pages, "done": 0, }): await finish_task(search_id) return try: await pixiv_search_live(url, pages, mode, phpsessid, search_id) await update_search_stats(search_id) await patch_explorer_cache(search_id, front=True) except TaskCanceled: pass except Exception as e: print(f"{search_id} failed: {e}") finally: await finish_task(search_id) async def deep_scan_search(search_id, phpsessid, deep_scan_types=None): allowed = exif_codes(deep_scan_types) if deep_scan_types is None: allowed = await get_search_deep_scan_types(search_id) targets = await get_deep_scan_targets(search_id, bool(allowed)) if not targets: return ACTIVE_TASKS[search_id].update( {"phase": "deep scanning", "total": len(targets), "done": 0, "raw_total": 0} ) await run_deep_scan( targets, phpsessid, task_id=search_id, allowed=allowed or None ) async def bg_user_task( search_id, user_id, phpsessid, deep_scan=True, deep_scan_types=None ): if not start_task(search_id, { "type": "user_search", "phase": "searching", "total": 1, "done": 0, }): await finish_task(search_id) return try: if task_canceled(search_id): raise TaskCanceled() result = await pixiv_user_posts(user_id, phpsessid) if task_canceled(search_id): raise TaskCanceled() post_ids = list(dict.fromkeys(result["post_ids"])) ACTIVE_TASKS[search_id]["done"] = 1 if not post_ids: return await save_search(search_id, post_ids, deep_scan, deep_scan_types) to_scan = await get_unscanned_post_ids(search_id) if to_scan: ACTIVE_TASKS[search_id].update({ "phase": "scanning", **await scan_progress_counts(search_id, len(to_scan)), }) await run_scan(to_scan, phpsessid, task_id=search_id, save_live=True) if deep_scan: await deep_scan_search(search_id, phpsessid, deep_scan_types) await cache_search_popularity(search_id) await update_search_stats(search_id) await patch_explorer_cache(search_id, front=True) except TaskCanceled: pass except Exception as e: print(f"{search_id} failed: {e}") finally: await finish_task(search_id) async def bg_user_batch_task(jobs, phpsessid, deep_scan=True, deep_scan_types=None): await asyncio.gather( *( bg_user_task( search_id, user_id, phpsessid, deep_scan, deep_scan_types ) for search_id, user_id in jobs ) ) async def bg_scan_task(search_id, post_ids, phpsessid, deep_scan=True): if not start_task(search_id, { "type": "scan", "phase": "scanning", **await scan_progress_counts(search_id, len(post_ids)), }): await finish_task(search_id) return completed = False try: await run_scan(post_ids, phpsessid, task_id=search_id, save_live=True) if deep_scan: await deep_scan_search(search_id, phpsessid) await cache_search_popularity(search_id) await update_search_stats(search_id) await patch_explorer_cache(search_id) completed = True except TaskCanceled: pass except Exception as e: print(f"{search_id} scan failed: {e}") finally: if not completed: await update_search_stats(search_id) await patch_explorer_cache(search_id) await finish_task(search_id) async def bg_search_and_scan_task( search_id, url, pages, mode, phpsessid, deep_scan=True, deep_scan_types=None ): if not start_task(search_id, { "type": "search+scan", "phase": "searching", "total": pages, "done": 0, }): await finish_task(search_id) return completed = False try: await pixiv_search_live( url, pages, mode, phpsessid, search_id, deep_scan, deep_scan_types ) to_scan = await get_unscanned_post_ids(search_id) if to_scan: ACTIVE_TASKS[search_id].update({ "phase": "scanning", **await scan_progress_counts(search_id, len(to_scan)), }) await run_scan(to_scan, phpsessid, task_id=search_id, save_live=True) if deep_scan: await deep_scan_search(search_id, phpsessid, deep_scan_types) await cache_search_popularity(search_id) await update_search_stats(search_id) await patch_explorer_cache(search_id, front=True) completed = True except TaskCanceled: pass except Exception as e: print(f"{search_id} failed: {e}") finally: if not completed: await update_search_stats(search_id) await patch_explorer_cache(search_id, front=True) await finish_task(search_id) async def bg_reclass_task( job_id, targets, filters, phpsessid, total=None, scanned=0, changes=None, updated=0, ): total = total or len(targets) changes = list(changes or []) if not start_task(job_id, { "type": "reclass", "phase": "scanning", "total": total, "done": scanned, "changed": len(changes), }): await finish_task(job_id) return try: for i in range(0, len(targets), RECLASS_BATCH_SIZE): if task_canceled(job_id): raise TaskCanceled() batch = targets[i : i + RECLASS_BATCH_SIZE] old = {item["post_id"]: item for item in batch} results = await run_scan( list(old), phpsessid, task_id=job_id, save_live=False, done_offset=scanned, ) if len(results) != len(batch): raise RuntimeError("scan batch incomplete") changed_results = [] for post_id, url, exif_type, deep_scanned in results: item = old.get(str(post_id)) if not item: continue short_url = url.replace(IMG_BASE, "", 1) if url else "" old_type = item["exif_type"] old_url = item["url"] if old_type != exif_type: changes.append( { "post_id": str(post_id), "from": exif_label(old_type), "to": exif_label(exif_type), } ) if old_type != exif_type or old_url != short_url: changed_results.append((post_id, url, exif_type, deep_scanned)) if changed_results: await save_scan_results(changed_results) await update_searches_for_post_ids( [str(result[0]) for result in changed_results] ) updated += len(changed_results) scanned += len(results) ACTIVE_TASKS[job_id].update({"done": scanned, "changed": len(changes)}) await finish_reclass_job(job_id, scanned, changes, updated, filters) except TaskCanceled: done = ACTIVE_TASKS.get(job_id, {}).get("done", 0) await finish_reclass_job(job_id, done, changes, updated, filters) except Exception as e: done = ACTIVE_TASKS.get(job_id, {}).get("done", 0) await finish_reclass_job(job_id, done, changes, updated, filters) print(f"{job_id} reclass failed: {e}") finally: await finish_task(job_id) class SearchRequest(BaseModel): url: str pages: int = 30 mode: str = "ai" action: str = "search" deep_scan: bool = True deep_scan_types: list = [] class UserSearchRequest(BaseModel): user_ids: list action: str = "search" deep_scan: bool = True deep_scan_types: list = [] class PostSearchRequest(BaseModel): post_ids: list class ScanRequest(BaseModel): search_id: str class RenameRequest(BaseModel): new_id: str class ReclassRequest(BaseModel): exif_types: list @app.on_event("startup") async def startup(): await asyncio.to_thread(cleanup_image_cache_root) if not TURSO_DB_URL: print("WARN: TURSO_DB_URL not set, skipping DB init") return try: await init_db() asyncio.create_task(warm_explorer_cache()) asyncio.create_task(resume_incomplete_scans()) asyncio.create_task(resume_incomplete_reclasses()) except Exception as e: print(f"WARN: DB init failed ({e}), will retry on first request") @app.on_event("shutdown") async def shutdown(): if TURSO_CLIENT and not TURSO_CLIENT.is_closed: await TURSO_CLIENT.aclose() if PIXIV_CLIENT and not PIXIV_CLIENT.closed: await PIXIV_CLIENT.close() async def backfill_search_stats(): try: resp = await turso_execute( [{"sql": "SELECT id FROM pi_searches WHERE total_searched = 0 AND post_ids != '[]' ORDER BY id DESC"}] ) rows = (resp.get("results") or [{}])[0].get("response", {}).get("result", {}).get("rows", []) ids = [row[0].get("value") for row in rows if row and row[0].get("value")] count = await update_search_stats_batch(ids) if count: print(f"Backfilled stats for {count} searches") except Exception as e: print(f"WARN: stats backfill failed: {e}") async def warm_explorer_cache(): try: await backfill_search_stats() await rebuild_explorer_cache() except Exception as e: print(f"WARN: explorer cache warm failed: {e}") async def resume_incomplete_scans(): try: search_ids = await get_incomplete_search_ids() started = 0 for search_id in search_ids: if search_id in ACTIVE_TASKS: continue post_ids = await get_unscanned_post_ids(search_id) deep_scan = await search_uses_deep_scan(search_id) if post_ids or deep_scan: asyncio.create_task( bg_scan_task(search_id, post_ids, PHPSESSID, deep_scan) ) started += 1 else: await update_search_stats(search_id) await patch_explorer_cache(search_id) if started: print(f"Resumed {started} incomplete scans") except Exception as e: print(f"WARN: scan resume failed: {e}") async def get_incomplete_reclass_jobs(): resp = await turso_execute( [ { "sql": "SELECT id, post_ids, filters, total, scanned, summary, log_enc " "FROM pi_reclass_jobs WHERE scanned < total ORDER BY id", } ] ) jobs = [] for row in first_rows(resp): job_id = cell_value(row, 0, "") post_ids = json.loads(cell_value(row, 1, "[]")) filters = json.loads(cell_value(row, 2, "[]")) total = int(cell_value(row, 3, "0")) scanned = min(int(cell_value(row, 4, "0")), total) summary = parse_json_maybe(cell_value(row, 5, "{}")) or {} changes = decrypt_log(cell_value(row, 6, "")) remaining = post_ids[scanned:] current = await get_scanned_post_ids(remaining) targets = [] for post_id in remaining: item = current.get(str(post_id), {}) targets.append( { "post_id": str(post_id), "url": item.get("url", ""), "exif_type": item.get("exif_type"), } ) jobs.append( { "id": job_id, "targets": targets, "filters": filters, "total": total or len(post_ids), "scanned": scanned, "changes": changes, "updated": int(summary.get("updated") or 0), } ) return jobs async def resume_incomplete_reclasses(): try: started = 0 for job in await get_incomplete_reclass_jobs(): job_id = job["id"] if job_id in ACTIVE_TASKS: continue if job["targets"]: asyncio.create_task( bg_reclass_task( job_id, job["targets"], job["filters"], PHPSESSID, total=job["total"], scanned=job["scanned"], changes=job["changes"], updated=job["updated"], ) ) started += 1 else: await finish_reclass_job( job_id, job["scanned"], job["changes"], job["updated"], job["filters"], ) if started: print(f"Resumed {started} incomplete reclasses") except Exception as e: print(f"WARN: reclass resume failed: {e}") @app.post("/api/submit") async def submit_search(req: SearchRequest, bg: BackgroundTasks): search_id = base26_time() phpsessid = PHPSESSID bg.add_task( bg_search_and_scan_task, search_id, req.url, req.pages, req.mode, phpsessid, req.deep_scan, req.deep_scan_types, ) return {"id": search_id, "status": "started"} @app.post("/api/submit_users") async def submit_users(req: UserSearchRequest, bg: BackgroundTasks): phpsessid = PHPSESSID user_ids = list(dict.fromkeys(int(u) for u in req.user_ids)) jobs = [(user_search_id(uid), uid) for uid in user_ids] bg.add_task( bg_user_batch_task, jobs, phpsessid, req.deep_scan, req.deep_scan_types, ) return {"ids": [search_id for search_id, _ in jobs], "status": "started"} @app.post("/api/post_search") async def post_search(req: PostSearchRequest): post_ids = list( dict.fromkeys(str(pid) for pid in req.post_ids if str(pid).isdigit()) ) if not post_ids: return {"error": "no post ids"} post_ids = post_ids[:POST_SEARCH_LIMIT] scanned = await get_scanned_post_ids(post_ids) result_pages = await get_result_pages(post_ids) items = [] for index, pid in enumerate(post_ids, 1): row = scanned.get(pid, {}) url = row.get("url", "") page_rows = result_pages.get(pid) or [] exif_types = sorted( {item["exif_type"] for item in page_rows if item["exif_type"]} ) if not exif_types and row.get("exif_type"): exif_types = [row["exif_type"]] pages = [item["page"] for item in page_rows] or [page_num_from_url(url)] items.append( { "post_id": pid, "scan_index": index, "url": url, "exif_type": row.get("exif_type"), "bookmark_count": row.get("bookmark_count"), "view_count": row.get("view_count"), "exif_types": exif_types, "scanned": pid in scanned, "pages": pages, "page_types": { str(item["page"]): item["exif_type"] or 0 for item in page_rows }, "page_count": len(pages), **image_links(pid, url), } ) return {"items": items, "total": len(items), "limit": POST_SEARCH_LIMIT} @app.post("/api/scan") async def scan_search(req: ScanRequest, bg: BackgroundTasks): phpsessid = PHPSESSID if req.search_id in ACTIVE_TASKS: return {"status": "active", **ACTIVE_TASKS[req.search_id]} counts = await get_search_counts(req.search_id) if not counts or not counts["raw_total"]: return {"error": "not found"} await turso_execute( [ { "sql": "UPDATE pi_searches SET deep_scan = 1 WHERE id = ?", "args": [{"type": "text", "value": req.search_id}], } ] ) to_scan = await get_unscanned_post_ids(req.search_id) deep_scan_types = await get_search_deep_scan_types(req.search_id) deep_targets = await get_deep_scan_targets( req.search_id, bool(deep_scan_types) ) popularity_ids = await get_missing_popularity_post_ids(req.search_id) if not to_scan and not deep_targets and not popularity_ids: return {"status": "already_scanned", "count": counts["raw_total"]} if not start_task(req.search_id, { "type": "scan", "phase": "queued", **await scan_progress_counts(req.search_id, len(to_scan)), }): return {"status": "canceling"} await publish_task_event(req.search_id) bg.add_task(bg_scan_task, req.search_id, to_scan, phpsessid, True) return {"status": "scanning", "total": counts["raw_total"], "to_scan": len(to_scan)} @app.post("/api/reclass") async def start_reclass(req: ReclassRequest, bg: BackgroundTasks): allowed = [ int(x) for x in req.exif_types if str(x).isdigit() and int(x) in (*EXIF_TYPE_TO_CODE.values(), 0) ] if not allowed: return {"error": "select at least one type"} job_id = base26_time() + "_reclass" targets = await get_reclass_targets(allowed) await save_reclass_job(job_id, [item["post_id"] for item in targets], allowed) if not targets: await finish_reclass_job(job_id, 0, [], 0, allowed) return {"id": job_id, "status": "empty"} ACTIVE_TASKS[job_id] = { "type": "reclass", "phase": "queued", "total": len(targets), "done": 0, } bg.add_task(bg_reclass_task, job_id, targets, allowed, PHPSESSID) return {"id": job_id, "status": "started", "total": len(targets)} @app.get("/api/reclass/{job_id}") async def get_reclass(job_id: str): resp = await turso_execute( [ { "sql": "SELECT id, post_ids, filters, total, scanned, changed, summary, log_enc, created_at " "FROM pi_reclass_jobs WHERE id = ?", "args": [{"type": "text", "value": job_id}], } ] ) rows = first_rows(resp) if not rows: return {"error": "not found"} row = rows[0] summary = parse_json_maybe(cell_value(row, 6, "{}")) or {} active = ACTIVE_TASKS.get(job_id) total = int(cell_value(row, 3, "0")) scanned = int(cell_value(row, 4, "0")) changed = int(cell_value(row, 5, "0")) if active: total = int(active.get("total") or total) scanned = int(active.get("done") or scanned) changed = int(active.get("changed") or changed) return { "id": cell_value(row, 0, ""), "kind": "reclass", "post_ids": json.loads(cell_value(row, 1, "[]")), "filters": json.loads(cell_value(row, 2, "[]")), "total": total, "scanned": scanned, "changed": changed, "summary": summary, "log": decrypt_log(cell_value(row, 7, "")), "created_at": int(cell_value(row, 8, "0")), "active": active, } @app.get("/api/searches") async def list_searches(page: int = 1): page = max(page, 1) if page == 1: cached = await cached_explorer_page() if cached: return cached return await rebuild_explorer_cache() return await query_explorer_page(page) @app.post("/api/searches/refresh") async def refresh_searches(): ids = await get_refresh_search_ids() count = await update_search_stats_batch(ids) await rebuild_explorer_cache() return { "status": "refreshed", "count": count, "more": len(ids) == REFRESH_SEARCH_LIMIT, } @app.get("/api/search/{search_id}") async def get_search(search_id: str): resp = await turso_execute( [ { "sql": "SELECT id, post_ids FROM pi_searches WHERE id = ?", "args": [{"type": "text", "value": search_id}], } ] ) results = resp.get("results") or [] if not results or "response" not in results[0]: return {"error": "not found"} rows = results[0]["response"].get("result", {}).get("rows", []) if not rows: return {"error": "not found"} row = rows[0] post_ids = json.loads(row[1].get("value", "[]")) scanned = await get_scanned_post_ids(post_ids) return { "id": row[0].get("value"), "post_ids": post_ids, "created_at": base26_to_time(row[0].get("value")), "scanned": scanned, } @app.get("/api/result_counts/{search_id}") async def result_counts(search_id: str): data = await get_result_counts(search_id) if not data or not data["raw_total"]: return {"error": "not found"} return {"search_id": search_id, **data} @app.get("/api/results/{search_id}") async def get_results( search_id: str, page: int = 1, exif_types: str = "", sort_bookmarks: bool = True, ): allowed = [ int(x) for x in exif_types.split(",") if x.isdigit() and int(x) in (*EXIF_TYPE_TO_CODE.values(), 0) ] await cache_search_popularity(search_id) data = await get_result_page( search_id, page, allowed if exif_types != "" else None, sort_bookmarks, ) if not data or not data["raw_total"]: return {"error": "not found"} return {"search_id": search_id, **data} @app.get("/api/results/{search_id}/download") async def download_results( search_id: str, background_tasks: BackgroundTasks, through: int, exif_types: str = "", ): if through < 1: raise HTTPException(status_code=400, detail="through must be positive") allowed = [ int(x) for x in exif_types.split(",") if x.isdigit() and int(x) in (*EXIF_TYPE_TO_CODE.values(), 0) ] targets = await get_download_targets( search_id, through, allowed if exif_types != "" else None ) if not targets: raise HTTPException(status_code=404, detail="no images found") path = await create_download_zip(targets) background_tasks.add_task(Path(path).unlink, missing_ok=True) name = "".join(c if c.isalnum() or c in "-_." else "_" for c in search_id) return FileResponse( path, media_type="application/zip", filename=f"{name.strip('.') or 'images'}.zip", background=background_tasks, ) @app.get("/api/i/{post_id}/t") async def get_image_thumb(post_id: str, p=0): p = int(p or 0) path = await create_webp(post_id, "", p, "t") return FileResponse( path, media_type="image/webp", headers={"Cache-Control": f"public, max-age={THUMB_MAX_AGE}"}, ) @app.get("/api/i/{post_id}/v") async def get_image_preview(post_id: str, p=0): p = int(p or 0) path = await create_webp(post_id, "", p, "v") return FileResponse( path, media_type="image/webp", headers={"Cache-Control": f"public, max-age={THUMB_MAX_AGE}"}, ) @app.get("/api/i/{post_id}/o") async def get_image_original(post_id: str, p=0): p = int(p or 0) image_url = await get_pixiv_image_url(post_id, p, "orig", PHPSESSID) data, content_type = await fetch_cached_image(post_id, p, image_url) filename = urlsplit(image_url).path.rsplit("/", 1)[-1] or f"{post_id}_p{p}.png" return Response( data, media_type=content_type, headers={ "Cache-Control": f"public, max-age={THUMB_MAX_AGE}", "Content-Disposition": f'attachment; filename="{filename}"', }, ) @app.get("/api/image/{post_id}/thumb") async def get_long_image_thumb(post_id: str, page: int = 0, p=None): return await get_image_thumb(post_id, page if p is None else p) @app.get("/api/image/{post_id}/full") async def get_long_image_full(post_id: str, page: int = 0, p=None): return await get_image_preview(post_id, page if p is None else p) @app.get("/api/thumb/{post_id}") async def get_thumb(post_id: str): return await get_image_thumb(post_id) @app.delete("/api/search/{search_id}") async def delete_search(search_id: str): active = ACTIVE_TASKS.get(search_id) cancel_task(search_id) if active: await publish_task_event(search_id) await delete_task_rows(search_id) await rebuild_explorer_cache() return {"status": "canceling" if active else "deleted"} @app.patch("/api/search/{search_id}") async def rename_search(search_id: str, req: RenameRequest): resp = await turso_execute( [ { "sql": "SELECT post_ids, found_exif, total_searched FROM pi_searches WHERE id = ?", "args": [{"type": "text", "value": search_id}], } ] ) results = resp.get("results") or [] if not results or "response" not in results[0]: return {"error": "not found"} rows = results[0]["response"].get("result", {}).get("rows", []) if not rows: return {"error": "not found"} post_ids_val = rows[0][0].get("value", "[]") found_exif = rows[0][1].get("value", "0") total_searched = rows[0][2].get("value", "0") await turso_execute( [ { "sql": "DELETE FROM pi_searches WHERE id = ?", "args": [{"type": "text", "value": search_id}], }, { "sql": "INSERT INTO pi_searches (id, post_ids, found_exif, total_searched) VALUES (?, ?, ?, ?)", "args": [ {"type": "text", "value": req.new_id}, {"type": "text", "value": post_ids_val}, {"type": "integer", "value": found_exif}, {"type": "integer", "value": total_searched}, ], }, ] ) await rebuild_explorer_cache() return {"status": "renamed", "new_id": req.new_id} @app.get("/api/progress") async def get_progress(): return [{"id": k, **v} for k, v in ACTIVE_TASKS.items()] @app.get("/api/events") async def events(request: Request): queue = asyncio.Queue(maxsize=16) TASK_EVENT_QUEUES.add(queue) async def stream(): try: while not await request.is_disconnected(): try: data = await asyncio.wait_for(queue.get(), timeout=25) yield f"data: {json.dumps(data)}\n\n" except asyncio.TimeoutError: yield ": keepalive\n\n" finally: TASK_EVENT_QUEUES.discard(queue) return StreamingResponse(stream(), media_type="text/event-stream") @app.get("/", response_class=HTMLResponse) @app.get("/index.html", response_class=HTMLResponse) def frontend(): directory = Path(FRONTEND_DIR) html = (directory / "index.html").read_text(encoding="utf-8") css = (directory / "style.css").read_text(encoding="utf-8") js = (directory / "app.js").read_text(encoding="utf-8") html = re.sub( r'', lambda _: f"", html, ) html = re.sub( r'', lambda _: f"", html, ) return HTMLResponse(html, headers={"Cache-Control": "no-cache"}) app.mount("/", StaticFiles(directory=FRONTEND_DIR, html=True), name="frontend")