From db8fc2eda5bcb709a65ca04a022a67884000e3df Mon Sep 17 00:00:00 2001 From: Zeecka Date: Sun, 14 Dec 2025 23:02:48 +0100 Subject: [PATCH] WIP --- aperisolve/analyzers/image_resize.py | 61 ++++++++++++---------------- aperisolve/analyzers/utils.py | 13 +----- aperisolve/config.py | 4 +- aperisolve/models.py | 23 ++++++++--- 4 files changed, 46 insertions(+), 55 deletions(-) diff --git a/aperisolve/analyzers/image_resize.py b/aperisolve/analyzers/image_resize.py index c294ad9..e975aac 100644 --- a/aperisolve/analyzers/image_resize.py +++ b/aperisolve/analyzers/image_resize.py @@ -4,7 +4,9 @@ import zlib import sqlite3 from pathlib import Path from typing import List -from .utils import update_data, DB_PATH +from .utils import update_data, unpack_ihdr +from ..app import app +from ..models import IHDR def write_recovered_png( @@ -17,6 +19,7 @@ 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") @@ -32,26 +35,18 @@ def write_recovered_png( out_f.write(full_png_data) -def calc_checksum( - header_chunk: bytearray, width_bytes: bytearray, height_bytes: bytearray -) -> bytearray: - """ - Calculates the CRC32 of the IHDR chunk with new dimensions. - header_chunk: The full IHDR chunk (Type + Data). - """ - # Reconstruct IHDR: [Type(4)] + [New Width(4)] + [New Height(4)] + [Rest of Data(5)] - new_header = header_chunk[:4] + width_bytes + height_bytes + header_chunk[12:] - - # Calculate CRC and return as 4 bytes (Big Endian) - return bytearray((zlib.crc32(new_header) & 0xFFFFFFFF).to_bytes(4, byteorder="big")) - - -def lookup_crc(crc_bytes: bytes) -> List[tuple[int, int]]: +def lookup_crc(crc_bytes: bytes) -> List[tuple[int, int, int, int, int]]: """ Queries the SQLite DB for the CRC and returns a list of (width, height) tuples. """ - # TODO query DB - return [] + with app.app_context(): + crc_b = int.from_bytes(crc_bytes, byteorder="big") + imgs = IHDR.query.filter_by(crc=crc_b).all() + results = [] + for img in imgs: + packed = img.packed + results.append(unpack_ihdr(packed)) + return results def analyze_image_resize(input_img: Path, output_dir: Path) -> None: @@ -91,30 +86,26 @@ def analyze_image_resize(input_img: Path, output_dir: Path) -> None: logs.append(f"Target CRC found: 0x{target_crc.hex()}") - candidates = [] saved_img_urls = [] - 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) + # 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 elt in db_matches: + w, h, _, _, _ = elt # TODO, use depth/color later if needed + 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)) - - else: - logs.append("Failure: No matching dimensions found in List or DB.") + # 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 f54042c..ff6e627 100644 --- a/aperisolve/analyzers/utils.py +++ b/aperisolve/analyzers/utils.py @@ -17,6 +17,7 @@ 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 + def update_data( output_dir: Path, new_data: dict[Any, Any], json_filename: str = "results.json" ) -> None: @@ -64,18 +65,6 @@ def pack_ihdr(width: int, height: int, depth: int, color: int, interlace: int) - color : 3 bits interlace : 1 bit """ - - 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") - return (width << 23) | (height << 7) | (depth << 4) | (color << 1) | interlace diff --git a/aperisolve/config.py b/aperisolve/config.py index dc12255..f3f0197 100644 --- a/aperisolve/config.py +++ b/aperisolve/config.py @@ -7,5 +7,5 @@ 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 +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 b44e73c..fc73a4f 100644 --- a/aperisolve/models.py +++ b/aperisolve/models.py @@ -59,9 +59,12 @@ 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[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: @@ -81,9 +84,15 @@ def create_crc_db() -> None: ) count = 0 - max_count = MAX_WIDTH * MAX_HEIGHT * len(bit_depth) * len(color_type) * len(interlace_method) + 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 + count += 1 # Construct IHDR chunk data ihdr_data = ( w_val.to_bytes(4, "big") @@ -102,7 +111,9 @@ def create_crc_db() -> None: db.session.add(ihdr_entry) if count % 10000 == 0: - print(f"Inserted {count}/{max_count} ({int((count*100)/max_count)}%) IHDR entries...") + print( + f"Inserted {count}/{max_count} ({int((count*100)/max_count)}%) IHDR entries..." + ) db.session.commit() # Commit all entries