diff --git a/Dockerfile b/Dockerfile index fc2198a..5e4eb5c 100644 --- a/Dockerfile +++ b/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"] \ No newline at end of file diff --git a/aperisolve/analyzers/image_resize.py b/aperisolve/analyzers/image_resize.py index 66a62c7..c294ad9 100644 --- a/aperisolve/analyzers/image_resize.py +++ b/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 - 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}") + # TODO query DB return [] @@ -252,59 +89,32 @@ 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: - # Create parent directories if they don't exist - output_dir.mkdir(parents=True, exist_ok=True) + + 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) - # Iterate through ALL candidates found - for w, h in candidates: - img_name = f"recovered_{w}x{h}.png" - output_path = output_dir / img_name - write_recovered_png(b, ihdr, w, h, output_path) + # Iterate through ALL candidates found + for w, h in candidates: + img_name = f"recovered_{w}x{h}.png" + output_path = output_dir / img_name + write_recovered_png(b, ihdr, w, h, output_path) - logs.append(f"Image saved: {img_name}") + logs.append(f"Image saved: {img_name}") - # Append the formatted URL to our list - saved_img_urls.append("/image/" + str(Path(output_dir.name) / img_name)) + # Append the formatted URL to our list + saved_img_urls.append("/image/" + str(Path(output_dir.name) / img_name)) - else: - logs.append("Failure: No matching dimensions found in List or DB.") + else: + logs.append("Failure: No matching dimensions found in List or DB.") # 4. Final Data Update output_data = { diff --git a/aperisolve/analyzers/utils.py b/aperisolve/analyzers/utils.py index c539abd..f54042c 100644 --- a/aperisolve/analyzers/utils.py +++ b/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: - """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 ( - crc INTEGER, - width INTEGER, - height INTEGER - ) +def pack_ihdr(width: int, height: int, depth: int, color: int, interlace: int) -> int: """ - ) + Pack IHDR fields into a single integer. - # Drop index if exists (rebuilding it at the end is faster) - cursor.execute("DROP INDEX IF EXISTS idx_crc") + Layout: + width : 16 bits + height : 16 bits + depth : 3 bits + color : 3 bits + interlace : 1 bit + """ - max_width = 2000 - max_height = 2000 - w = range(1, max_width + 1) - h = range(1, max_height + 1) + if not (0 < width <= 0xFFFF): + 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") - # 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] + return (width << 23) | (height << 7) | (depth << 4) | (color << 1) | interlace - # Create the generator - all_combinations = itertools.product( - w, h, bit_depth, color_type, compression_method, filter_method, interlace_method - ) - batch_size = 100000 - batch = [] - count = 0 +def unpack_ihdr(packed: int) -> tuple[int, int, int, int, int]: + """ + Unpack packed IHDR integer into components: + (width, height, depth, color, interlace) + """ - for w_val, h_val, bd, ct, comp, filt, inter in all_combinations: - # Construct IHDR chunk data - ihdr_data = ( - w_val.to_bytes(4, "big") - + h_val.to_bytes(4, "big") - + bytes([bd, ct, comp, filt, inter]) - ) + interlace = packed & 0b1 + color = (packed >> 1) & 0b111 + depth = (packed >> 4) & 0b111 + height = (packed >> 7) & 0xFFFF + width = (packed >> 23) & 0xFFFF - # 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.") + return width, height, depth, color, interlace diff --git a/aperisolve/config.py b/aperisolve/config.py index fe000d3..dc12255 100644 --- a/aperisolve/config.py +++ b/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 \ No newline at end of file diff --git a/aperisolve/models.py b/aperisolve/models.py index 730bc63..b44e73c 100644 --- a/aperisolve/models.py +++ b/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.")