venus/aperisolve
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

CRC Lookup on PostgreSQL WIP

Browse Source
This commit is contained in:
Zeecka
2025-12-14 22:32:14 +01:00
parent 50578eca12
commit d3928427e5
5 changed files with 117 additions and 298 deletions
Download Patch File Download Diff File Expand all files Collapse all files
aperisolve/analyzers/image_resize.py
+19 -209
View File
@@ -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 = {
aperisolve/analyzers/utils.py
+31 -85
View File
@@ -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
aperisolve/config.py
+3
View File
@@ -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
aperisolve/models.py
+62 -2
View File
@@ -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.")