diff --git a/aperisolve/analyzers/image_resize.py b/aperisolve/analyzers/image_resize.py index e975aac..7fdfcb4 100644 --- a/aperisolve/analyzers/image_resize.py +++ b/aperisolve/analyzers/image_resize.py @@ -1,17 +1,15 @@ """Image size bruteforcer for common PNG sizes.""" -import zlib -import sqlite3 from pathlib import Path from typing import List -from .utils import update_data, unpack_ihdr + from ..app import app from ..models import IHDR +from .utils import PNG, update_data def write_recovered_png( png_bytes: bytearray, - ihdr_offset: int, width: int, height: int, output_path: Path, @@ -19,40 +17,102 @@ def write_recovered_png( """ Rebuilds a PNG by patching IHDR width/height and writes it to disk. """ - # TODO, use depth/color later if needed height_bytes = height.to_bytes(4, byteorder="big") width_bytes = width.to_bytes(4, byteorder="big") - # Splicing: [Start...IHDR+4] + [W] + [H] + [IHDR+12...End] full_png_data = ( - png_bytes[: ihdr_offset + 4] + png_bytes[:16] # Splicing: [Start...IHDR+4] + [W] + [H] + [IHDR+12...End] + width_bytes + height_bytes - + png_bytes[ihdr_offset + 12 :] + + png_bytes[24:] ) with output_path.open("wb") as out_f: out_f.write(full_png_data) -def lookup_crc(crc_bytes: bytes) -> List[tuple[int, int, int, int, int]]: +def lookup_crc(png: PNG) -> List[PNG]: """ - Queries the SQLite DB for the CRC and returns a list of (width, height) tuples. + Look up PNG images in the database by their CRC value. + + Args: + png (PNG): A PNG object containing the CRC value to search for. + + Returns: + List[PNG]: A list of PNG objects with matching CRC values, reconstructed from their packed + representations. """ with app.app_context(): - crc_b = int.from_bytes(crc_bytes, byteorder="big") - imgs = IHDR.query.filter_by(crc=crc_b).all() + imgs = IHDR.query.filter_by(crc=png.crc).all() results = [] for img in imgs: packed = img.packed - results.append(unpack_ihdr(packed)) + results.append(PNG.from_packed(packed, crc=png.crc)) + return results + + +def search_height_crc(png: PNG) -> List[PNG]: + """ + Search for PNG images with matching CRC by iterating through possible heights. + + This function attempts to find alternative PNG configurations that produce the same + CRC checksum as the input PNG by systematically varying the height parameter while + keeping all other properties constant. + + Args: + png (PNG): The reference PNG object whose CRC checksum will be matched. + + Returns: + List[PNG]: A list of PNG objects with matching CRC values. Returns an empty list + if no matches are found. + + Note: + The function iterates through heights from 100 to 3500 (inclusive). + This may be computationally expensive for large ranges. + """ + results = [] + for height in range(100, 3501): + candidate = PNG( + width=png.width, + height=height, + bit_depth=png.bit_depth, + color_type=png.color_type, + interlace=png.interlace, + ) + if candidate.crc == png.crc: + results.append(candidate) return results def analyze_image_resize(input_img: Path, output_dir: Path) -> None: """ - Analyze an image submission using python-based resize bruteforce. - Strategy: Check common sizes first, then fall back to DB lookup. + Analyze a PNG image and recover resized versions based on CRC matching. + + This function attempts to find matching image dimensions for a given PNG file + by looking up its CRC (Cyclic Redundancy Check) in a database. If matches are + found, it recovers and saves PNG images with the matched dimensions. + + Args: + input_img (Path): Path to the input PNG image file to analyze. + output_dir (Path): Path to the directory where recovered PNG images will be saved. + + Returns: + None + + Raises: + Handles exceptions gracefully by updating output data with error status. + + Side Effects: + - Creates output directory if it doesn't exist. + - Saves recovered PNG images to the output directory. + - Updates data in the output directory with analysis results via update_data(). + + Notes: + - If the PNG has invalid structure or IHDR chunk is not first, an error is logged. + - The function performs a two-step lookup: first by CRC, then by height/CRC if no matches + found. + - Recovered images are named as "recovered_{width}x{height}.png". + - Analysis results (logs, status, and image URLs) are persisted via update_data(). """ # 1. Setup Paths @@ -64,48 +124,48 @@ def analyze_image_resize(input_img: Path, output_dir: Path) -> None: try: with input_img.open("rb") as image: - f = image.read() + img_bytes = bytearray(image.read()) + png = PNG.from_bytearray(img_bytes) - b = bytearray(f) - - # Find IHDR start - ihdr = b.find(b"\x49\x48\x44\x52") - if ihdr == -1: - logs.append("Failure: PNG header (IHDR) not found.") + if png is None: + logs.append( + "Failure: IHDR chunk is not the first chunk, or PNG has invalid structure." + ) update_data( output_dir, - {"image_resize": {"status": "error", "error": "PNG header not found."}}, + { + "image_resize": { + "status": "error", + "error": "IHDR chunk is not the first chunk, or PNG has invalid structure.", + } + }, ) return - # Calculate Chunk Length (The 4 bytes BEFORE IHDR tag) - chunk_length = int.from_bytes((b[ihdr - 4 : ihdr]), byteorder="big") + 4 - - # Extract Target CRC (The 4 bytes AFTER the chunk data) - target_crc = b[ihdr + chunk_length : ihdr + chunk_length + 4] - - logs.append(f"Target CRC found: 0x{target_crc.hex()}") + logs.append(f"Target CRC found: 0x{png.crc:08x}") saved_img_urls = [] - db_matches = lookup_crc(target_crc) - if db_matches: + db_matches = lookup_crc(png) + if not db_matches: + db_matches = search_height_crc(png) + if not db_matches: + logs.append("Failure: No matching dimensions found.") + else: # Create parent directories if they don't exist output_dir.mkdir(parents=True, exist_ok=True) # Iterate through ALL candidates found - for elt in db_matches: - w, h, _, _, _ = elt # TODO, use depth/color later if needed - img_name = f"recovered_{w}x{h}.png" + for im_match in db_matches: + img_name = f"recovered_{im_match.width}x{im_match.height}.png" output_path = output_dir / img_name - write_recovered_png(b, ihdr, w, h, output_path) + write_recovered_png( + img_bytes, im_match.width, im_match.height, output_path + ) 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)) - 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 ff6e627..13d7196 100644 --- a/aperisolve/analyzers/utils.py +++ b/aperisolve/analyzers/utils.py @@ -2,13 +2,14 @@ import fcntl import json +import math import os +import struct import threading -from pathlib import Path -from typing import Any -import itertools import zlib -import sqlite3 +from functools import cached_property +from pathlib import Path +from typing import Any, Self _thread_lock = threading.Lock() @@ -54,30 +55,218 @@ def update_data( fcntl.flock(lock, fcntl.LOCK_UN) -def pack_ihdr(width: int, height: int, depth: int, color: int, interlace: int) -> int: - """ - Pack IHDR fields into a single integer. - - Layout: - width : 16 bits - height : 16 bits - depth : 3 bits - color : 3 bits - interlace : 1 bit - """ - return (width << 23) | (height << 7) | (depth << 4) | (color << 1) | interlace +def get_resolutions(): + BASE_WIDTHS = ( + list(range(16, 257, 16)) # 16 → 256 + + list(range(320, 1025, 32)) # 320 → 1024 + + list(range(1280, 2561, 64)) # 1280 → 2560 + + list(range(3000, 4097, 128)) # 3K–4K + + list(range(5120, 8193, 256)) # 5K–8K + + [10000] # upper bound + ) + ASPECT_RATIOS = [ + # Screens / digital + (1, 1), # square + (4, 3), + (3, 2), + (16, 10), + (16, 9), + (21, 9), + # Photography / print + (5, 4), # 8x10 + (7, 5), # 5x7 + (3, 2), # 4x6 (already included, kept for clarity) + (2, 3), # poster + # Paper standards + (1414, 1000), # ISO A-series (A4, A3, A2, ...) + (11, 85), # US Letter + (14, 85), # Legal + (17, 11), # Tabloid + ] + resolutions = set() + for w in BASE_WIDTHS: + for ar_w, ar_h in ASPECT_RATIOS: + h = math.floor(round(w * ar_h / ar_w)) + if 1 <= h <= 10000: + resolutions.add((w, h)) + resolutions.add((h, w)) # portrait / landscape + return sorted(resolutions) -def unpack_ihdr(packed: int) -> tuple[int, int, int, int, int]: +def get_valid_depth_color_pairs(): + VALID_COLOR_DEPTHS = { + 0: [1, 2, 4, 8, 16], + 2: [8, 16], + 3: [1, 2, 4, 8], + 4: [8, 16], + 6: [8, 16], + } + for color, depths in VALID_COLOR_DEPTHS.items(): + for depth in depths: + yield depth, color + + +class PNG: """ - Unpack packed IHDR integer into components: - (width, height, depth, color, interlace) + A class representing a PNG image header (IHDR chunk). + This class encapsulates the metadata from a PNG file's IHDR (Image Header) chunk, + including image dimensions, bit depth, color type, interlacing method, and CRC checksum. + It provides methods for parsing PNG files, packing/unpacking header data, and generating + PNG crc if omitted. """ - interlace = packed & 0b1 - color = (packed >> 1) & 0b111 - depth = (packed >> 4) & 0b111 - height = (packed >> 7) & 0xFFFF - width = (packed >> 23) & 0xFFFF + width: int + height: int + bit_depth: int + color_type: int + crc: int + interlace: int = 0 - return width, height, depth, color, interlace + def __init__( + self, + width: int, + height: int, + bit_depth: int, + color_type: int, + crc: int | None = None, + interlace: int = 0, + ): + """Initialize a PNG IHDR (Image Header) chunk. + + Args: + width (int): The width of the image in pixels. + height (int): The height of the image in pixels. + bit_depth (int): The number of bits per sample or per palette index (1, 2, 4, 8, or 16). + color_type (int): The color type of the image (0=Grayscale, 2=RGB, 3=Indexed, + 4=Grayscale+Alpha, 6=RGBA). + crc (int | None, optional): CRC (Cyclic Redundancy Check) checksum for the chunk. + If None, it will be calculated automatically. Defaults + to None. + interlace (int, optional): The interlace method (0=no interlace, 1=Adam7 interlace). + Defaults to 0. + """ + self.width = width + self.height = height + self.bit_depth = bit_depth + self.color_type = color_type + self.crc = ( + crc + or zlib.crc32( + b"IHDR" + + width.to_bytes(4, "big") + + height.to_bytes(4, "big") + + bytes((bit_depth, color_type, 0, 0, interlace)) + ) + & 0xFFFFFFFF + ) + self.interlace = interlace + + @classmethod + def from_packed(cls, packed: int, crc: int | None = None) -> Self: + """Create a PNG instance from a packed integer representation. + + Parse a packed integer containing width, height, bit depth, color type and interlace + bits and return a PNG instance. The CRC may be supplied; if omitted it will be + computed by the constructor. + + Args: + packed (int): Packed integer with fields laid out as described in the `packed` property. + crc (int | None): Optional CRC value for the IHDR chunk. + + Returns: + Self: PNG instance with parsed fields. + """ + interlace = packed & 0b1 + color = (packed >> 1) & 0b111 + depth = (packed >> 4) & 0b111 + height = (packed >> 7) & 0xFFFF + width = (packed >> 23) & 0xFFFF + return cls( + width=width, + height=height, + bit_depth=depth, + color_type=color, + crc=crc, + interlace=interlace, + ) + + @cached_property + def packed(self) -> int: + """Pack IHDR fields into a single integer. + + Returns: + int: Packed integer encoding the IHDR fields. + + Layout: + width : 16 bits + height : 16 bits + depth : 3 bits + color : 3 bits + interlace : 1 bit + """ + return ( + ((self.width & 0xFFFF) << 23) + | ((self.height & 0xFFFF) << 7) + | ((self.bit_depth & 0b111) << 4) + | ((self.color_type & 0b111) << 1) + | (self.interlace & 0b1) + ) + + @classmethod + def from_bytearray(cls, data: bytearray) -> Self: + """Create a PNG image instance from a bytearray containing PNG file data. + + Parses the PNG signature and IHDR (Image Header) chunk to extract image metadata + including dimensions, bit depth, color type, and interlacing information. + + Args: + data (bytearray): A bytearray containing the complete PNG file data, starting + with the PNG signature (8 bytes: \\x89PNG\\r\\n\\x1a\\n). + + Returns: + Self | None: A new instance of the PNG image class with parsed IHDR metadata + (width, height, bit_depth, color_type, crc, interlace) if the data is + valid and contains a proper IHDR chunk. Returns None if the PNG signature + is invalid or the IHDR chunk type is not found. + + Raises: + struct.error: If the data is too short or malformed for unpacking IHDR fields. + IndexError: If the bytearray is shorter than expected for the PNG structure. + """ + + PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n" + if data[:8] != PNG_SIGNATURE: + raise ValueError("Invalid PNG signature") + offset = 8 # Read IHDR chunk + length = struct.unpack(">I", data[offset : offset + 4])[0] + offset += 4 + chunk_type = data[offset : offset + 4] + offset += 4 + if chunk_type != b"IHDR": + raise ValueError("IHDR chunk not found at expected position") + + ihdr_data = data[offset : offset + length] + offset += length + + crc = struct.unpack(">I", data[offset : offset + 4])[0] + + # Unpack IHDR fields + width, height, depth, color, _, _, interlace = struct.unpack( + ">IIBBBBB", ihdr_data + ) + + return cls( + width=width, + height=height, + bit_depth=depth, + color_type=color, + crc=crc, + interlace=interlace, + ) + + def __repr__(self) -> str: + return ( + f"PNG(width={self.width}, height={self.height}, " + f"bit_depth={self.bit_depth}, color_type={self.color_type}, " + f"interlace={self.interlace}, crc=0x{self.crc:08x})" + ) diff --git a/aperisolve/cleanup.py b/aperisolve/cleanup.py index 130542d..0e59175 100644 --- a/aperisolve/cleanup.py +++ b/aperisolve/cleanup.py @@ -14,7 +14,28 @@ MAX_STORE_TIME = int(getenv("MAX_STORE_TIME", "259200")) # 3 days by default def cleanup_old_entries() -> None: - """Clean up old entries in the database and file system.""" + """ + Clean up old and incomplete entries from the database and file system. + + This function performs two main cleanup operations: + + 1. Submission cleanup: + - Deletes submissions with "pending" or "running" status that have exceeded + the maximum allowed processing time (MAX_PENDING_TIME). + - Deletes completed submissions ("done" status) that have missing or buggy + results (results.json file not found). + + 2. Image cleanup: + - Deletes images older than MAX_STORE_TIME along with all associated + submissions and their result folders. + - Deletes orphaned images (with no submissions) older than MAX_PENDING_TIME + and removes their associated result folders from the file system. + + All database changes are committed at the end of the operation. + + Returns: + None + """ now = time.time() for submission in Submission.query.all(): # type: ignore if ( diff --git a/aperisolve/config.py b/aperisolve/config.py index f3f0197..fe000d3 100644 --- a/aperisolve/config.py +++ b/aperisolve/config.py @@ -6,6 +6,3 @@ 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 = 5 # Used for IHDR lookup table -MAX_HEIGHT = 5 # Used for IHDR lookup table diff --git a/aperisolve/models.py b/aperisolve/models.py index fc73a4f..1f2d897 100644 --- a/aperisolve/models.py +++ b/aperisolve/models.py @@ -1,16 +1,14 @@ -"""This module defines the database models for the Aperisolve application.""" +"""This module defines the database models for the Aperi'Solve 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, BigInteger +from sqlalchemy import BigInteger, Boolean, Column, DateTime, Float, Integer, String + +from aperisolve.analyzers.utils import PNG, get_resolutions, get_valid_depth_color_pairs db: SQLAlchemy = SQLAlchemy() @@ -60,68 +58,76 @@ class Submission(db.Model): # type: ignore 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) + iid = Column(Integer, primary_key=True, autoincrement=True) + crc = Column(BigInteger, nullable=False) + packed = Column(BigInteger, nullable=False) def create_crc_db() -> None: - w = range(1, MAX_WIDTH + 1) - h = range(1, MAX_HEIGHT + 1) + """ + Creates a database of common PNG IHDR (Image Header) entries by generating + combinations of standard PNG parameters such as resolutions, bit depths, + color types, and interlace methods. The function checks for existing entries + in the database to avoid duplicates and inserts new entries in batches of + 5000, committing the changes to the database periodically. + Returns: + None + """ # 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] + resolutions = get_resolutions() + bit_color_pairs = get_valid_depth_color_pairs() + interlace_methods = [0, 1] - # Create the generator - all_combinations = itertools.product( - w, h, bit_depth, color_type, compression_method, filter_method, interlace_method + combinations = itertools.product( + resolutions, + bit_color_pairs, + interlace_methods, ) 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 + with db.session.no_autoflush: # pylint: disable=no-member + for (w, h), (bd, ct), inter in combinations: + p = PNG(width=w, height=h, bit_depth=bd, color_type=ct, interlace=inter) - # Compute packed - packed = pack_ihdr(w_val, h_val, bd, ct, inter) + exists = db.session.execute( # pylint: disable=no-member + db.select(IHDR.iid).where(IHDR.packed == p.packed) + ).first() - # Create and add IHDR entry to database - ihdr_entry = IHDR(crc=crc, packed=packed) - db.session.add(ihdr_entry) + if exists: # Skip if already present + continue - if count % 10000 == 0: - print( - f"Inserted {count}/{max_count} ({int((count*100)/max_count)}%) IHDR entries..." - ) - db.session.commit() + db.session.add(IHDR(packed=p.packed, crc=p.crc)) # pylint: disable=no-member + count += 1 - # Commit all entries - db.session.commit() + if count % 5000 == 0: + print(f"Inserted {count} common IHDR entries...") + db.session.commit() + + db.session.commit() # pylint: disable=no-member + print(f"Precomputed {count} common IHDR entries.") def init_db(app: Flask) -> None: - """Initialize the database with the given Flask app.""" + """ + Initialize the database by creating tables and populating lookup data. + + This function creates the database schema if it doesn't already exist, + and populates the CRC lookup table for PNG IHDR validation. + + Args: + app (Flask): The Flask application instance used to establish the application context + for database operations. + + Returns: + None + + Side Effects: + - Creates all database tables if the "image" table doesn't exist + - Populates the IHDR CRC lookup table with initial data + - Prints status messages to console during initialization + """ with app.app_context(): if not db.engine.dialect.has_table(db.engine.connect(), "image"): print("Creating database...") diff --git a/aperisolve/workers.py b/aperisolve/workers.py index 956c328..acb6b8e 100644 --- a/aperisolve/workers.py +++ b/aperisolve/workers.py @@ -8,19 +8,43 @@ from .analyzers.binwalk import analyze_binwalk from .analyzers.decomposer import analyze_decomposer from .analyzers.exiftool import analyze_exiftool from .analyzers.foremost import analyze_foremost +from .analyzers.image_resize import analyze_image_resize from .analyzers.outguess import analyze_outguess +from .analyzers.pngcheck import analyze_pngcheck from .analyzers.steghide import analyze_steghide from .analyzers.strings import analyze_strings from .analyzers.zsteg import analyze_zsteg -from .analyzers.pngcheck import analyze_pngcheck -from .analyzers.image_resize import analyze_image_resize from .app import app, db from .config import RESULT_FOLDER from .models import Image, Submission def analyze_image(submission_hash: str) -> None: - """Analyze an image submission in separate threads.""" + """ + Analyze an image submission by running multiple analysis tools concurrently. + + This function retrieves a submission and its associated image from the database, + then executes various image analysis tools in parallel threads. The analysis tools + include binwalk, decomposer, exiftool, foremost, strings, pngcheck, steghide, + zsteg, and image resize. If deep analysis is enabled, additional tools like + outguess are included. + + Args: + submission_hash (str): The unique hash identifier of the submission to analyze. + + Returns: + None + + Raises: + Implicitly catches and handles exceptions during analysis, setting submission + status to "error" if an exception occurs. + + Side Effects: + - Queries the database for the submission and associated image + - Updates submission status in the database ("running" -> "completed" or "error") + - Creates result directories and generates analysis output files + - Modifies the database session and commits changes + """ with app.app_context(): submission: Submission = Submission.query.get(submission_hash) # type: ignore image = Image.query.get_or_404(submission.image_hash) # type: ignore @@ -29,7 +53,7 @@ def analyze_image(submission_hash: str) -> None: return submission.status = "running" # type: ignore - db.session.commit() + db.session.commit() # pylint: disable=no-member try: img_path: Path = Path(image.file) @@ -77,4 +101,4 @@ def analyze_image(submission_hash: str) -> None: except Exception: submission.status = "error" finally: - db.session.commit() + db.session.commit() # pylint: disable=no-member