CRC Lookup on PostgreSQL WIP
This commit is contained in:
Zeecka
+2
-2
@@ -23,6 +23,6 @@ RUN mkdir -p /data
|
||||
|
||||
RUN pip install --no-cache-dir -r /aperisolve/requirements.txt
|
||||
|
||||
RUN PYTHONPATH=/ python -c "from aperisolve.analyzers.utils import create_crc_db; create_crc_db()"
|
||||
ENV PYTHONUNBUFFERED=1
|
||||
|
||||
CMD ["gunicorn", "-w", "4", "-b", "0.0.0.0:5000", "aperisolve.wsgi:application"]
|
||||
CMD ["gunicorn", "-w", "4", "-b", "0.0.0.0:5000", "--access-logfile", "-", "--error-logfile", "-", "--log-level", "info", "--capture-output", "aperisolve.wsgi:application"]
|
||||
@@ -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)
|
||||
|
||||
# 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)
|
||||
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)
|
||||
|
||||
logs.append(f"Image saved: {img_name}")
|
||||
# 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)
|
||||
|
||||
# Append the formatted URL to our list
|
||||
saved_img_urls.append("/image/" + str(Path(output_dir.name) / img_name))
|
||||
logs.append(f"Image saved: {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 = {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
+62
-2
@@ -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.")
|
||||
|
||||
Reference in New Issue
Block a user