This commit is contained in:
Zeecka
2025-12-14 23:02:48 +01:00
parent d3928427e5
commit db8fc2eda5
4 changed files with 46 additions and 55 deletions
+15 -24
View File
@@ -4,7 +4,9 @@ import zlib
import sqlite3 import sqlite3
from pathlib import Path from pathlib import Path
from typing import List 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( 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. 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") height_bytes = height.to_bytes(4, byteorder="big")
width_bytes = width.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) 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. 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: def analyze_image_resize(input_img: Path, output_dir: Path) -> None:
@@ -91,19 +86,16 @@ def analyze_image_resize(input_img: Path, output_dir: Path) -> None:
logs.append(f"Target CRC found: 0x{target_crc.hex()}") logs.append(f"Target CRC found: 0x{target_crc.hex()}")
candidates = []
saved_img_urls = [] saved_img_urls = []
db_matches = lookup_crc(target_crc) db_matches = lookup_crc(target_crc)
if db_matches: if db_matches:
candidates.extend(db_matches)
if candidates:
# Create parent directories if they don't exist # Create parent directories if they don't exist
output_dir.mkdir(parents=True, exist_ok=True) 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:
w, h, _, _, _ = elt # TODO, use depth/color later if needed
img_name = f"recovered_{w}x{h}.png" img_name = f"recovered_{w}x{h}.png"
output_path = output_dir / img_name output_path = output_dir / img_name
write_recovered_png(b, ihdr, w, h, output_path) write_recovered_png(b, ihdr, w, h, output_path)
@@ -112,7 +104,6 @@ def analyze_image_resize(input_img: Path, output_dir: Path) -> None:
# 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.")
+1 -12
View File
@@ -17,6 +17,7 @@ MAX_PENDING_TIME = int(os.getenv("MAX_PENDING_TIME", "600")) # 10 minutes by de
# PATH CONFIGURATION # PATH CONFIGURATION
# Prioritize /data for Docker persistence, fallback to local directory for bare metal # Prioritize /data for Docker persistence, fallback to local directory for bare metal
def update_data( def update_data(
output_dir: Path, new_data: dict[Any, Any], json_filename: str = "results.json" output_dir: Path, new_data: dict[Any, Any], json_filename: str = "results.json"
) -> None: ) -> None:
@@ -64,18 +65,6 @@ def pack_ihdr(width: int, height: int, depth: int, color: int, interlace: int) -
color : 3 bits color : 3 bits
interlace : 1 bit 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 return (width << 23) | (height << 7) | (depth << 4) | (color << 1) | interlace
+2 -2
View File
@@ -7,5 +7,5 @@ WORKER_FILES = ["binwalk", "foremost", "steghide", "zsteg", "image_resize"]
RESULT_FOLDER = Path(__file__).parent.resolve() / "results" 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
+14 -3
View File
@@ -59,7 +59,10 @@ class Submission(db.Model): # type: ignore
class IHDR(db.Model): # type: ignore class IHDR(db.Model): # type: ignore
"""IHDR CRC lookup table""" """IHDR CRC lookup table"""
iid : Column[int] = Column(Integer, primary_key=True, autoincrement=True, unique=True)
iid: Column[int] = Column(
Integer, primary_key=True, autoincrement=True, unique=True
)
crc: Column[int] = Column(BigInteger) crc: Column[int] = Column(BigInteger)
packed: Column[int] = Column(Integer) packed: Column[int] = Column(Integer)
@@ -81,7 +84,13 @@ def create_crc_db() -> None:
) )
count = 0 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: for w_val, h_val, bd, ct, comp, filt, inter in all_combinations:
count += 1 count += 1
# Construct IHDR chunk data # Construct IHDR chunk data
@@ -102,7 +111,9 @@ def create_crc_db() -> None:
db.session.add(ihdr_entry) db.session.add(ihdr_entry)
if count % 10000 == 0: 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() db.session.commit()
# Commit all entries # Commit all entries