CRC Lookup on PostgreSQL WIP

Dockerfile

@@ -23,6 +23,6 @@ RUN mkdir -p /data
 
 RUN pip install --no-cache-dir -r /aperisolve/requirements.txt
 
-RUN PYTHONPATH=/ python -c "from aperisolve.analyzers.utils import create_crc_db; create_crc_db()"
+ENV PYTHONUNBUFFERED=1
 
-CMD ["gunicorn", "-w", "4", "-b", "0.0.0.0:5000", "aperisolve.wsgi:application"]
+CMD ["gunicorn", "-w", "4", "-b", "0.0.0.0:5000", "--access-logfile", "-", "--error-logfile", "-", "--log-level", "info", "--capture-output", "aperisolve.wsgi:application"]

aperisolve/analyzers/image_resize.py

@@ -7,144 +7,6 @@ from typing import List
 from .utils import update_data, DB_PATH
 
 
-# Full list of expected sizes
-EXPECTED_SIZES = [
-    # -------------------------
-    # ICONS (Square)
-    # -------------------------
-    (16, 16),
-    (20, 20),
-    (24, 24),
-    (28, 28),
-    (32, 32),
-    (48, 48),
-    (64, 64),
-    (72, 72),
-    (96, 96),
-    (128, 128),
-    (192, 192),
-    (256, 256),
-    (384, 384),
-    (512, 512),
-    (768, 768),
-    (1024, 1024),
-    (1536, 1536),
-    (2048, 2048),
-    (4096, 4096),
-    # -------------------------
-    # 1:1 Square (General)
-    # -------------------------
-    (300, 300),
-    (500, 500),
-    (800, 800),
-    (1000, 1000),
-    (1500, 1500),
-    (2000, 2000),
-    (2500, 2500),
-    (3000, 3000),
-    # -------------------------
-    # 4:3 Aspect (Classic Screens)
-    # -------------------------
-    (320, 240),
-    (640, 480),
-    (800, 600),
-    (960, 720),
-    (1024, 768),
-    (1280, 960),
-    (1600, 1200),
-    (2048, 1536),
-    (2560, 1920),
-    (3200, 2400),
-    (4096, 3072),
-    # -------------------------
-    # 3:2 Aspect (Photography)
-    # -------------------------
-    (600, 400),
-    (900, 600),
-    (1200, 800),
-    (1500, 1000),
-    (1800, 1200),
-    (2400, 1600),
-    (3000, 2000),
-    (3600, 2400),
-    (4500, 3000),
-    (6000, 4000),
-    # -------------------------
-    # 16:9 Aspect (Video / Web)
-    # -------------------------
-    (426, 240),
-    (640, 360),
-    (854, 480),
-    (960, 540),
-    (1280, 720),
-    (1366, 768),
-    (1600, 900),
-    (1920, 1080),
-    (2560, 1440),
-    (3200, 1800),
-    (3440, 1935),
-    (3840, 2160),
-    (5120, 2880),
-    (7680, 4320),
-    (10240, 5760),
-    # -------------------------
-    # 16:10 Aspect (Laptops)
-    # -------------------------
-    (1280, 800),
-    (1440, 900),
-    (1680, 1050),
-    (1920, 1200),
-    (2560, 1600),
-    (2880, 1800),
-    (3840, 2400),
-    (5120, 3200),
-    # -------------------------
-    # Mobile Portrait (9:16)
-    # -------------------------
-    (720, 1280),
-    (1080, 1920),
-    (1440, 2560),
-    (2160, 3840),
-    (1440, 2960),
-    (1080, 2400),
-    (1170, 2532),
-    (1284, 2778),
-    (1290, 2796),
-    # -------------------------
-    # Social Media Sizes
-    # -------------------------
-    # Instagram
-    (1080, 1080),
-    (1080, 1350),
-    (1080, 1920),
-    # Facebook
-    (1200, 628),
-    (1200, 1200),
-    # Twitter/X
-    (1200, 675),
-    (1500, 500),
-    # YouTube Thumbnails
-    (1280, 720),
-    (1920, 1080),
-    # -------------------------
-    # Banners / Headers
-    # -------------------------
-    (1600, 500),
-    (1920, 500),
-    (2560, 700),
-    (3000, 1000),
-    (3840, 1200),
-    # -------------------------
-    # Print-like Digital Sizes
-    # -------------------------
-    (2550, 3300),  # Letter @ 300 DPI
-    (2480, 3508),  # A4 @ 300 DPI
-    (3508, 4961),  # A3 @ 300 DPI
-    (4961, 7016),  # A2 @ 300 DPI
-    (7016, 9933),  # A1 @ 300 DPI
-]
-
-
 def write_recovered_png(
     png_bytes: bytearray,
     ihdr_offset: int,
@@ -184,36 +46,11 @@ def calc_checksum(
     return bytearray((zlib.crc32(new_header) & 0xFFFFFFFF).to_bytes(4, byteorder="big"))
 
 
-def lookup_crc(crc_bytes: bytes, logs: List[str]) -> List[tuple[int, int]]:
+def lookup_crc(crc_bytes: bytes) -> List[tuple[int, int]]:
     """
     Queries the SQLite DB for the CRC and returns a list of (width, height) tuples.
     """
-    # Locate DB relative to this script file
+    # TODO query DB
-    db_path = Path(DB_PATH)
-
-    if not db_path.exists():
-        logs.append(f"Error: Database not found at {db_path}")
-        return []
-
-    try:
-        target_crc = int.from_bytes(crc_bytes, byteorder="big")
-    except ValueError:
-        logs.append(f"Invalid CRC bytes: {crc_bytes!r}")
-        return []
-
-    conn = sqlite3.connect(db_path)
-    cursor = conn.cursor()
-
-    # Query the database
-    cursor.execute("SELECT width, height FROM ihdr WHERE crc = ?", (target_crc,))
-    results = cursor.fetchall()  # Returns list of (width, height)
-    conn.close()
-
-    if results:
-        logs.append(f"Database: Found {len(results)} match(es) for CRC {target_crc}")
-        return results
-
-    logs.append(f"Database: No match found for CRC {target_crc}")
     return []
 
 
@@ -252,43 +89,16 @@ def analyze_image_resize(input_img: Path, output_dir: Path) -> None:
         # Extract Target CRC (The 4 bytes AFTER the chunk data)
         target_crc = b[ihdr + chunk_length : ihdr + chunk_length + 4]
 
-        # Isolate the header chunk (Type + Data)
-        header_chunk = b[ihdr : ihdr + chunk_length]
-
         logs.append(f"Target CRC found: 0x{target_crc.hex()}")
 
-        match_found = False
         candidates = []
-
-        # --- STRATEGY 1: Check Common Sizes ---
-        for size in EXPECTED_SIZES:
-            width, height = size
-
-            # Convert dimensions to 4-byte Big Endian arrays
-            width_bytes = bytearray(width.to_bytes(4, byteorder="big"))
-            height_bytes = bytearray(height.to_bytes(4, byteorder="big"))
-
-            # Check if this size matches the file's CRC
-            if target_crc == calc_checksum(header_chunk, width_bytes, height_bytes):
-                match_found = True
-                logs.append(
-                    f"Success (Common List): Match found! Dimensions: {width}x{height}"
-                )
-                candidates.append((width, height))
-                break  # Stop after first match
-
-        # --- STRATEGY 2: Check SQLite DB ---
-        if not match_found:
-            logs.append("No common size matched. Checking database...")
-            db_matches = lookup_crc(target_crc, logs)
-            if db_matches:
-                candidates.extend(db_matches)
-                match_found = True
-
-        # --- Generate Output Images ---
         saved_img_urls = []
 
-        if match_found and candidates:
+    
+        db_matches = lookup_crc(target_crc)
+        if db_matches:
+            candidates.extend(db_matches)
+            if candidates:
                 # Create parent directories if they don't exist
                 output_dir.mkdir(parents=True, exist_ok=True)
 

aperisolve/analyzers/utils.py

@@ -16,12 +16,6 @@ MAX_PENDING_TIME = int(os.getenv("MAX_PENDING_TIME", "600"))  # 10 minutes by de
 
 # PATH CONFIGURATION
 # Prioritize /data for Docker persistence, fallback to local directory for bare metal
-DB_NAME = "ihdr_crcs.db"
-if os.path.exists("/data"):
-    DB_PATH = os.path.join("/data", DB_NAME)
-else:
-    DB_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), DB_NAME)
-
 
 def update_data(
     output_dir: Path, new_data: dict[Any, Any], json_filename: str = "results.json"
@@ -59,90 +53,42 @@ def update_data(
                 fcntl.flock(lock, fcntl.LOCK_UN)
 
 
-def create_crc_db() -> None:
+def pack_ihdr(width: int, height: int, depth: int, color: int, interlace: int) -> int:
-    """Creates and populates CRC lookup database with CRC, Width and Heigth tuples"""
-    if os.path.exists(DB_PATH):
-        print(f"Database already exists at {DB_PATH}. Skipping generation.")
-        return
-
-    # Use the global DB_PATH variable so it writes to /data/ in Docker
-    # and /analyzers/ in local dev
-    print(f"Generating database at: {DB_PATH}")
-    conn = sqlite3.connect(DB_PATH)
-
-    cursor = conn.cursor()
-
-    # Optimization: Turn off safety checks for bulk insertion speed
-    cursor.execute("PRAGMA synchronous = OFF")
-    cursor.execute("PRAGMA journal_mode = MEMORY")
-
-    # Create table
-    # We store the CRC as an integer for faster indexing/storage
-    cursor.execute(
     """
-        CREATE TABLE IF NOT EXISTS ihdr (
+    Pack IHDR fields into a single integer.
-            crc INTEGER,
+
-            width INTEGER,
+    Layout:
-            height INTEGER
+    width     : 16 bits
-        )
+    height    : 16 bits
+    depth     : 3 bits
+    color     : 3 bits
+    interlace : 1 bit
     """
-    )
 
-    # Drop index if exists (rebuilding it at the end is faster)
+    if not (0 < width <= 0xFFFF):
-    cursor.execute("DROP INDEX IF EXISTS idx_crc")
+        raise ValueError("width out of range")
+    if not (0 < height <= 0xFFFF):
+        raise ValueError("height out of range")
+    if depth not in (1, 2, 4, 8, 16):
+        raise ValueError("invalid depth")
+    if color not in (0, 2, 3, 4, 6):
+        raise ValueError("invalid color")
+    if interlace not in (0, 1):
+        raise ValueError("invalid interlace")
 
-    max_width = 2000
+    return (width << 23) | (height << 7) | (depth << 4) | (color << 1) | interlace
-    max_height = 2000
-    w = range(1, max_width + 1)
-    h = range(1, max_height + 1)
 
-    # Standard PNG IHDR parameters
-    bit_depth = [1, 2, 4, 8, 16]
-    color_type = [0, 2, 3, 4, 6]
-    compression_method = [0]  # This is always 0 according to the PNG spec
-    filter_method = [0]  # This is always 0 according to the PNG spec
-    interlace_method = [0, 1]
 
-    # Create the generator
+def unpack_ihdr(packed: int) -> tuple[int, int, int, int, int]:
-    all_combinations = itertools.product(
+    """
-        w, h, bit_depth, color_type, compression_method, filter_method, interlace_method
+    Unpack packed IHDR integer into components:
-    )
+    (width, height, depth, color, interlace)
+    """
 
-    batch_size = 100000
+    interlace = packed & 0b1
-    batch = []
+    color = (packed >> 1) & 0b111
-    count = 0
+    depth = (packed >> 4) & 0b111
+    height = (packed >> 7) & 0xFFFF
+    width = (packed >> 23) & 0xFFFF
 
-    for w_val, h_val, bd, ct, comp, filt, inter in all_combinations:
+    return width, height, depth, color, interlace
-        # Construct IHDR chunk data
-        ihdr_data = (
-            w_val.to_bytes(4, "big")
-            + h_val.to_bytes(4, "big")
-            + bytes([bd, ct, comp, filt, inter])
-        )
-
-        # Calculate CRC
-        crc = zlib.crc32(b"IHDR" + ihdr_data) & 0xFFFFFFFF
-
-        # Add to batch
-        batch.append((crc, w_val, h_val))
-
-        if len(batch) >= batch_size:
-            cursor.executemany("INSERT INTO ihdr VALUES (?,?,?)", batch)
-            conn.commit()
-            batch = []
-            count += batch_size
-            if count % 1000000 == 0:
-                print(f"Inserted {count:,} rows...")
-
-    # Insert remaining records
-    if batch:
-        cursor.executemany("INSERT INTO ihdr VALUES (?,?,?)", batch)
-        conn.commit()
-
-    print("Insertion complete. Creating Index for faster lookups...")
-    # Creating the index allows for faster lookups
-    cursor.execute("CREATE INDEX idx_crc ON ihdr(crc)")
-    conn.commit()
-
-    conn.close()
-    print("Database ready.")

aperisolve/config.py

@@ -6,3 +6,6 @@ IMAGE_EXTENSIONS = [".png", ".jpg", ".jpeg", ".gif", ".bmp", ".webp", ".tiff"]
 WORKER_FILES = ["binwalk", "foremost", "steghide", "zsteg", "image_resize"]
 
 RESULT_FOLDER = Path(__file__).parent.resolve() / "results"
+
+MAX_WIDTH = 2000  # Used for IHDR lookup table
+MAX_HEIGHT = 2000  # Used for IHDR lookup table

aperisolve/models.py

@@ -1,11 +1,16 @@
 """This module defines the database models for the Aperisolve application."""
 
+import zlib
+import itertools
+
 from datetime import datetime, timezone
 from decimal import Decimal
 
+from aperisolve.analyzers.utils import pack_ihdr
+from aperisolve.config import MAX_HEIGHT, MAX_WIDTH
 from flask import Flask
 from flask_sqlalchemy import SQLAlchemy
-from sqlalchemy import Boolean, Column, DateTime, Float, Integer, String
+from sqlalchemy import Boolean, Column, DateTime, Float, Integer, String, BigInteger
 
 db: SQLAlchemy = SQLAlchemy()
 
@@ -52,10 +57,65 @@ class Submission(db.Model):  # type: ignore
     image_hash = Column(String, db.ForeignKey("image.hash"), nullable=False)
 
 
+class IHDR(db.Model):  # type: ignore
+    """IHDR CRC lookup table"""
+    iid : Column[int] = Column(Integer, primary_key=True, autoincrement=True, unique=True)
+    crc : Column[int] = Column(BigInteger)
+    packed : Column[int] = Column(Integer)
+
+
+def create_crc_db() -> None:
+    w = range(1, MAX_WIDTH + 1)
+    h = range(1, MAX_HEIGHT + 1)
+
+    # Standard PNG IHDR parameters
+    bit_depth = [1, 2, 4, 8, 16]
+    color_type = [0, 2, 3, 4, 6]
+    compression_method = [0]  # This is always 0 according to the PNG spec
+    filter_method = [0]  # This is always 0 according to the PNG spec
+    interlace_method = [0, 1]
+
+    # Create the generator
+    all_combinations = itertools.product(
+        w, h, bit_depth, color_type, compression_method, filter_method, interlace_method
+    )
+
+    count = 0
+    max_count = MAX_WIDTH * MAX_HEIGHT * len(bit_depth) * len(color_type) * len(interlace_method)
+    for w_val, h_val, bd, ct, comp, filt, inter in all_combinations:
+        count +=1
+        # Construct IHDR chunk data
+        ihdr_data = (
+            w_val.to_bytes(4, "big")
+            + h_val.to_bytes(4, "big")
+            + bytes([bd, ct, comp, filt, inter])
+        )
+
+        # Calculate CRC
+        crc = zlib.crc32(b"IHDR" + ihdr_data) & 0xFFFFFFFF
+
+        # Compute packed
+        packed = pack_ihdr(w_val, h_val, bd, ct, inter)
+
+        # Create and add IHDR entry to database
+        ihdr_entry = IHDR(crc=crc, packed=packed)
+        db.session.add(ihdr_entry)
+
+        if count % 10000 == 0:
+            print(f"Inserted {count}/{max_count} ({int((count*100)/max_count)}%) IHDR entries...")
+            db.session.commit()
+
+    # Commit all entries
+    db.session.commit()
+
+
 def init_db(app: Flask) -> None:
     """Initialize the database with the given Flask app."""
     with app.app_context():
         if not db.engine.dialect.has_table(db.engine.connect(), "image"):
             print("Creating database...")
             db.create_all()
-            print("Database created successfully.")
+            print("Database structure created successfully.")
+            print("Filling IHDR lookup table...")
+            create_crc_db()
+            print("IHDR table filled successfully.")