WIP

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(
+def lookup_crc(crc_bytes: bytes) -> List[tuple[int, int, int, int, int]]:
-    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]]:
     """
     Queries the SQLite DB for the CRC and returns a list of (width, height) tuples.
     """
-    # TODO query DB
+    with app.app_context():
-    return []
+        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)
+            # Create parent directories if they don't exist
-            if candidates:
+            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
+            # Iterate through ALL candidates found
-                for w, h in candidates:
+            for elt in db_matches:
-                    img_name = f"recovered_{w}x{h}.png"
+                w, h, _, _, _ = elt  # TODO, use depth/color later if needed
-                    output_path = output_dir / img_name
+                img_name = f"recovered_{w}x{h}.png"
-                    write_recovered_png(b, ihdr, w, h, output_path)
+                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
+                # Append the formatted URL to our list
-                    saved_img_urls.append("/image/" + str(Path(output_dir.name) / img_name))
+                saved_img_urls.append("/image/" + str(Path(output_dir.name) / img_name))
-
+        else:
-            else:
+            logs.append("Failure: No matching dimensions found in List or DB.")
-                logs.append("Failure: No matching dimensions found in List or DB.")
 
         # 4. Final Data Update
         output_data = {

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
 
 

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_WIDTH = 5  # Used for IHDR lookup table
-MAX_HEIGHT = 2000  # Used for IHDR lookup table
+MAX_HEIGHT = 5  # Used for IHDR lookup table

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)
+    iid: Column[int] = Column(
-    packed : Column[int] = Column(Integer)
+        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