Stealth

API/app.py

@@ -2,6 +2,9 @@ from fastapi import FastAPI, Query, BackgroundTasks
 from fastapi.responses import FileResponse
 import aiohttp
 import asyncio
+import gzip
+import io
+import json
 import time
 import tempfile
 import zipfile
@@ -9,6 +12,9 @@ import os
 from pydantic import BaseModel
 from typing import List, Dict, Optional
 
+import numpy as np
+from PIL import Image
+
 
 img_base = 'https://i.pximg.net/img-original/img/'
 COMMENTS_ROOTS_URL = "https://www.pixiv.net/ajax/illusts/comments/roots"
@@ -108,6 +114,57 @@ def parse_png_metadata(data):
         index += chunk_len + 4
     return None
 
+def byteize(alpha):
+    alpha = alpha.T.reshape((-1,))
+    alpha = alpha[:(alpha.shape[0] // 8) * 8]
+    alpha = np.bitwise_and(alpha, 1)
+    alpha = alpha.reshape((-1, 8))
+    alpha = np.packbits(alpha, axis=1)
+    return alpha
+
+class LSBExtractor:
+    def __init__(self, alpha):
+        self.data = byteize(alpha)
+        self.pos = 0
+
+    def get_next_n_bytes(self, n):
+        n_bytes = self.data[self.pos:self.pos + n]
+        self.pos += n
+        return bytearray(n_bytes)
+
+    def read_32bit_integer(self):
+        bytes_list = self.get_next_n_bytes(4)
+        if len(bytes_list) == 4:
+            return int.from_bytes(bytes_list, byteorder="big")
+        return None
+
+def extract_stealth_metadata(image: Image.Image) -> dict:
+    if "A" not in image.getbands():
+        raise AssertionError("image format")
+    alpha = np.array(image.getchannel("A"))
+    reader = LSBExtractor(alpha)
+    magic = "stealth_pngcomp"
+    read_magic = reader.get_next_n_bytes(len(magic)).decode("utf-8")
+    if magic != read_magic:
+        raise AssertionError("magic number")
+    read_len = reader.read_32bit_integer()
+    if read_len is None:
+        raise AssertionError("length missing")
+    read_len = read_len // 8
+    json_data = reader.get_next_n_bytes(read_len)
+    json_data = json.loads(gzip.decompress(json_data).decode("utf-8"))
+    if "Comment" in json_data and isinstance(json_data["Comment"], str):
+        json_data["Comment"] = json.loads(json_data["Comment"])
+    return json_data
+
+def has_stealth_png_bytes(data: bytes) -> bool:
+    try:
+        image = Image.open(io.BytesIO(data))
+        extract_stealth_metadata(image)
+        return True
+    except Exception:
+        return False
+
 def format_comment_text(comment):
     text = comment.get("comment") or ""
     if not text and comment.get("stampId"):
@@ -193,6 +250,10 @@ async def fetch_post_comments(post_id, session, limit, semaphore):
 
     return str(post_id), results
 
+async def fetch_image_data(session, url):
+    async with session.get(url) as response:
+        return await response.read()
+
 async def process_post(post_id, session, semaphore):
     async with semaphore:
         try:
@@ -217,21 +278,37 @@ async def process_post(post_id, session, semaphore):
 
             for s, e in chunks:
                 chunk_tasks = [get_exif(image_urls[i], session) for i in range(s, e)]
-                results = await asyncio.gather(*chunk_tasks)
+                results = await asyncio.gather(*chunk_tasks, return_exceptions=True)
 
                 for image_url, metadata in zip(image_urls[s:e], results):
-                    exif_type = determine_exif_type(metadata)
+                    if isinstance(metadata, Exception):
+                        exif_type = None
+                    else:
+                        exif_type = determine_exif_type(metadata)
                     if exif_type not in ['photoshop', 'celsys', None]:
                         return post_id, image_url
+
+            for s, e in chunks:
+                stealth_tasks = [fetch_image_data(session, image_urls[i]) for i in range(s, e)]
+                stealth_results = await asyncio.gather(*stealth_tasks, return_exceptions=True)
+
+                for image_url, data in zip(image_urls[s:e], stealth_results):
+                    if isinstance(data, Exception) or data is None:
+                        continue
+                    if has_stealth_png_bytes(data):
+                        return post_id, image_url
                     
             return post_id, None
         except Exception:
             return post_id, None
 
 async def fetch_image_bytes(session, url, post_id, semaphore):
-    async with semaphore:
-        async with session.get(url) as response:
-            return post_id, await response.read()
+    if semaphore:
+        async with semaphore:
+            async with session.get(url) as response:
+                return post_id, await response.read()
+    async with session.get(url) as response:
+        return post_id, await response.read()
 
 @app.get("/allimages")
 async def all_images(

Dockerfile

@@ -8,7 +8,7 @@ ENV PYTHONDONTWRITEBYTECODE=1
 ENV PYTHONUNBUFFERED=1
 WORKDIR /app
 
-RUN pip install --no-cache-dir fastapi aiohttp uvicorn pydantic requests
+RUN pip install --no-cache-dir fastapi aiohttp uvicorn pydantic requests numpy pillow
 
 COPY --chown=user ./API /app
-CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]