Image resize and linting
This commit is contained in:
@@ -1,17 +1,15 @@
|
|||||||
"""Image size bruteforcer for common PNG sizes."""
|
"""Image size bruteforcer for common PNG sizes."""
|
||||||
|
|
||||||
import zlib
|
|
||||||
import sqlite3
|
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import List
|
from typing import List
|
||||||
from .utils import update_data, unpack_ihdr
|
|
||||||
from ..app import app
|
from ..app import app
|
||||||
from ..models import IHDR
|
from ..models import IHDR
|
||||||
|
from .utils import PNG, update_data
|
||||||
|
|
||||||
|
|
||||||
def write_recovered_png(
|
def write_recovered_png(
|
||||||
png_bytes: bytearray,
|
png_bytes: bytearray,
|
||||||
ihdr_offset: int,
|
|
||||||
width: int,
|
width: int,
|
||||||
height: int,
|
height: int,
|
||||||
output_path: Path,
|
output_path: Path,
|
||||||
@@ -19,40 +17,102 @@ 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")
|
||||||
|
|
||||||
# Splicing: [Start...IHDR+4] + [W] + [H] + [IHDR+12...End]
|
|
||||||
full_png_data = (
|
full_png_data = (
|
||||||
png_bytes[: ihdr_offset + 4]
|
png_bytes[:16] # Splicing: [Start...IHDR+4] + [W] + [H] + [IHDR+12...End]
|
||||||
+ width_bytes
|
+ width_bytes
|
||||||
+ height_bytes
|
+ height_bytes
|
||||||
+ png_bytes[ihdr_offset + 12 :]
|
+ png_bytes[24:]
|
||||||
)
|
)
|
||||||
|
|
||||||
with output_path.open("wb") as out_f:
|
with output_path.open("wb") as out_f:
|
||||||
out_f.write(full_png_data)
|
out_f.write(full_png_data)
|
||||||
|
|
||||||
|
|
||||||
def lookup_crc(crc_bytes: bytes) -> List[tuple[int, int, int, int, int]]:
|
def lookup_crc(png: PNG) -> List[PNG]:
|
||||||
"""
|
"""
|
||||||
Queries the SQLite DB for the CRC and returns a list of (width, height) tuples.
|
Look up PNG images in the database by their CRC value.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
png (PNG): A PNG object containing the CRC value to search for.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
List[PNG]: A list of PNG objects with matching CRC values, reconstructed from their packed
|
||||||
|
representations.
|
||||||
"""
|
"""
|
||||||
with app.app_context():
|
with app.app_context():
|
||||||
crc_b = int.from_bytes(crc_bytes, byteorder="big")
|
imgs = IHDR.query.filter_by(crc=png.crc).all()
|
||||||
imgs = IHDR.query.filter_by(crc=crc_b).all()
|
|
||||||
results = []
|
results = []
|
||||||
for img in imgs:
|
for img in imgs:
|
||||||
packed = img.packed
|
packed = img.packed
|
||||||
results.append(unpack_ihdr(packed))
|
results.append(PNG.from_packed(packed, crc=png.crc))
|
||||||
|
return results
|
||||||
|
|
||||||
|
|
||||||
|
def search_height_crc(png: PNG) -> List[PNG]:
|
||||||
|
"""
|
||||||
|
Search for PNG images with matching CRC by iterating through possible heights.
|
||||||
|
|
||||||
|
This function attempts to find alternative PNG configurations that produce the same
|
||||||
|
CRC checksum as the input PNG by systematically varying the height parameter while
|
||||||
|
keeping all other properties constant.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
png (PNG): The reference PNG object whose CRC checksum will be matched.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
List[PNG]: A list of PNG objects with matching CRC values. Returns an empty list
|
||||||
|
if no matches are found.
|
||||||
|
|
||||||
|
Note:
|
||||||
|
The function iterates through heights from 100 to 3500 (inclusive).
|
||||||
|
This may be computationally expensive for large ranges.
|
||||||
|
"""
|
||||||
|
results = []
|
||||||
|
for height in range(100, 3501):
|
||||||
|
candidate = PNG(
|
||||||
|
width=png.width,
|
||||||
|
height=height,
|
||||||
|
bit_depth=png.bit_depth,
|
||||||
|
color_type=png.color_type,
|
||||||
|
interlace=png.interlace,
|
||||||
|
)
|
||||||
|
if candidate.crc == png.crc:
|
||||||
|
results.append(candidate)
|
||||||
return results
|
return results
|
||||||
|
|
||||||
|
|
||||||
def analyze_image_resize(input_img: Path, output_dir: Path) -> None:
|
def analyze_image_resize(input_img: Path, output_dir: Path) -> None:
|
||||||
"""
|
"""
|
||||||
Analyze an image submission using python-based resize bruteforce.
|
Analyze a PNG image and recover resized versions based on CRC matching.
|
||||||
Strategy: Check common sizes first, then fall back to DB lookup.
|
|
||||||
|
This function attempts to find matching image dimensions for a given PNG file
|
||||||
|
by looking up its CRC (Cyclic Redundancy Check) in a database. If matches are
|
||||||
|
found, it recovers and saves PNG images with the matched dimensions.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
input_img (Path): Path to the input PNG image file to analyze.
|
||||||
|
output_dir (Path): Path to the directory where recovered PNG images will be saved.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
None
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
Handles exceptions gracefully by updating output data with error status.
|
||||||
|
|
||||||
|
Side Effects:
|
||||||
|
- Creates output directory if it doesn't exist.
|
||||||
|
- Saves recovered PNG images to the output directory.
|
||||||
|
- Updates data in the output directory with analysis results via update_data().
|
||||||
|
|
||||||
|
Notes:
|
||||||
|
- If the PNG has invalid structure or IHDR chunk is not first, an error is logged.
|
||||||
|
- The function performs a two-step lookup: first by CRC, then by height/CRC if no matches
|
||||||
|
found.
|
||||||
|
- Recovered images are named as "recovered_{width}x{height}.png".
|
||||||
|
- Analysis results (logs, status, and image URLs) are persisted via update_data().
|
||||||
"""
|
"""
|
||||||
|
|
||||||
# 1. Setup Paths
|
# 1. Setup Paths
|
||||||
@@ -64,48 +124,48 @@ def analyze_image_resize(input_img: Path, output_dir: Path) -> None:
|
|||||||
|
|
||||||
try:
|
try:
|
||||||
with input_img.open("rb") as image:
|
with input_img.open("rb") as image:
|
||||||
f = image.read()
|
img_bytes = bytearray(image.read())
|
||||||
|
png = PNG.from_bytearray(img_bytes)
|
||||||
|
|
||||||
b = bytearray(f)
|
if png is None:
|
||||||
|
logs.append(
|
||||||
# Find IHDR start
|
"Failure: IHDR chunk is not the first chunk, or PNG has invalid structure."
|
||||||
ihdr = b.find(b"\x49\x48\x44\x52")
|
)
|
||||||
if ihdr == -1:
|
|
||||||
logs.append("Failure: PNG header (IHDR) not found.")
|
|
||||||
update_data(
|
update_data(
|
||||||
output_dir,
|
output_dir,
|
||||||
{"image_resize": {"status": "error", "error": "PNG header not found."}},
|
{
|
||||||
|
"image_resize": {
|
||||||
|
"status": "error",
|
||||||
|
"error": "IHDR chunk is not the first chunk, or PNG has invalid structure.",
|
||||||
|
}
|
||||||
|
},
|
||||||
)
|
)
|
||||||
return
|
return
|
||||||
|
|
||||||
# Calculate Chunk Length (The 4 bytes BEFORE IHDR tag)
|
logs.append(f"Target CRC found: 0x{png.crc:08x}")
|
||||||
chunk_length = int.from_bytes((b[ihdr - 4 : ihdr]), byteorder="big") + 4
|
|
||||||
|
|
||||||
# Extract Target CRC (The 4 bytes AFTER the chunk data)
|
|
||||||
target_crc = b[ihdr + chunk_length : ihdr + chunk_length + 4]
|
|
||||||
|
|
||||||
logs.append(f"Target CRC found: 0x{target_crc.hex()}")
|
|
||||||
|
|
||||||
saved_img_urls = []
|
saved_img_urls = []
|
||||||
|
|
||||||
db_matches = lookup_crc(target_crc)
|
db_matches = lookup_crc(png)
|
||||||
if db_matches:
|
if not db_matches:
|
||||||
|
db_matches = search_height_crc(png)
|
||||||
|
if not db_matches:
|
||||||
|
logs.append("Failure: No matching dimensions found.")
|
||||||
|
else:
|
||||||
# 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 elt in db_matches:
|
for im_match in db_matches:
|
||||||
w, h, _, _, _ = elt # TODO, use depth/color later if needed
|
img_name = f"recovered_{im_match.width}x{im_match.height}.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(
|
||||||
|
img_bytes, im_match.width, im_match.height, 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:
|
|
||||||
logs.append("Failure: No matching dimensions found in List or DB.")
|
|
||||||
|
|
||||||
# 4. Final Data Update
|
# 4. Final Data Update
|
||||||
output_data = {
|
output_data = {
|
||||||
|
|||||||
+214
-25
@@ -2,13 +2,14 @@
|
|||||||
|
|
||||||
import fcntl
|
import fcntl
|
||||||
import json
|
import json
|
||||||
|
import math
|
||||||
import os
|
import os
|
||||||
|
import struct
|
||||||
import threading
|
import threading
|
||||||
from pathlib import Path
|
|
||||||
from typing import Any
|
|
||||||
import itertools
|
|
||||||
import zlib
|
import zlib
|
||||||
import sqlite3
|
from functools import cached_property
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Any, Self
|
||||||
|
|
||||||
_thread_lock = threading.Lock()
|
_thread_lock = threading.Lock()
|
||||||
|
|
||||||
@@ -54,30 +55,218 @@ def update_data(
|
|||||||
fcntl.flock(lock, fcntl.LOCK_UN)
|
fcntl.flock(lock, fcntl.LOCK_UN)
|
||||||
|
|
||||||
|
|
||||||
def pack_ihdr(width: int, height: int, depth: int, color: int, interlace: int) -> int:
|
def get_resolutions():
|
||||||
"""
|
BASE_WIDTHS = (
|
||||||
Pack IHDR fields into a single integer.
|
list(range(16, 257, 16)) # 16 → 256
|
||||||
|
+ list(range(320, 1025, 32)) # 320 → 1024
|
||||||
Layout:
|
+ list(range(1280, 2561, 64)) # 1280 → 2560
|
||||||
width : 16 bits
|
+ list(range(3000, 4097, 128)) # 3K–4K
|
||||||
height : 16 bits
|
+ list(range(5120, 8193, 256)) # 5K–8K
|
||||||
depth : 3 bits
|
+ [10000] # upper bound
|
||||||
color : 3 bits
|
)
|
||||||
interlace : 1 bit
|
ASPECT_RATIOS = [
|
||||||
"""
|
# Screens / digital
|
||||||
return (width << 23) | (height << 7) | (depth << 4) | (color << 1) | interlace
|
(1, 1), # square
|
||||||
|
(4, 3),
|
||||||
|
(3, 2),
|
||||||
|
(16, 10),
|
||||||
|
(16, 9),
|
||||||
|
(21, 9),
|
||||||
|
# Photography / print
|
||||||
|
(5, 4), # 8x10
|
||||||
|
(7, 5), # 5x7
|
||||||
|
(3, 2), # 4x6 (already included, kept for clarity)
|
||||||
|
(2, 3), # poster
|
||||||
|
# Paper standards
|
||||||
|
(1414, 1000), # ISO A-series (A4, A3, A2, ...)
|
||||||
|
(11, 85), # US Letter
|
||||||
|
(14, 85), # Legal
|
||||||
|
(17, 11), # Tabloid
|
||||||
|
]
|
||||||
|
resolutions = set()
|
||||||
|
for w in BASE_WIDTHS:
|
||||||
|
for ar_w, ar_h in ASPECT_RATIOS:
|
||||||
|
h = math.floor(round(w * ar_h / ar_w))
|
||||||
|
if 1 <= h <= 10000:
|
||||||
|
resolutions.add((w, h))
|
||||||
|
resolutions.add((h, w)) # portrait / landscape
|
||||||
|
return sorted(resolutions)
|
||||||
|
|
||||||
|
|
||||||
def unpack_ihdr(packed: int) -> tuple[int, int, int, int, int]:
|
def get_valid_depth_color_pairs():
|
||||||
|
VALID_COLOR_DEPTHS = {
|
||||||
|
0: [1, 2, 4, 8, 16],
|
||||||
|
2: [8, 16],
|
||||||
|
3: [1, 2, 4, 8],
|
||||||
|
4: [8, 16],
|
||||||
|
6: [8, 16],
|
||||||
|
}
|
||||||
|
for color, depths in VALID_COLOR_DEPTHS.items():
|
||||||
|
for depth in depths:
|
||||||
|
yield depth, color
|
||||||
|
|
||||||
|
|
||||||
|
class PNG:
|
||||||
"""
|
"""
|
||||||
Unpack packed IHDR integer into components:
|
A class representing a PNG image header (IHDR chunk).
|
||||||
(width, height, depth, color, interlace)
|
This class encapsulates the metadata from a PNG file's IHDR (Image Header) chunk,
|
||||||
|
including image dimensions, bit depth, color type, interlacing method, and CRC checksum.
|
||||||
|
It provides methods for parsing PNG files, packing/unpacking header data, and generating
|
||||||
|
PNG crc if omitted.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
interlace = packed & 0b1
|
width: int
|
||||||
color = (packed >> 1) & 0b111
|
height: int
|
||||||
depth = (packed >> 4) & 0b111
|
bit_depth: int
|
||||||
height = (packed >> 7) & 0xFFFF
|
color_type: int
|
||||||
width = (packed >> 23) & 0xFFFF
|
crc: int
|
||||||
|
interlace: int = 0
|
||||||
|
|
||||||
return width, height, depth, color, interlace
|
def __init__(
|
||||||
|
self,
|
||||||
|
width: int,
|
||||||
|
height: int,
|
||||||
|
bit_depth: int,
|
||||||
|
color_type: int,
|
||||||
|
crc: int | None = None,
|
||||||
|
interlace: int = 0,
|
||||||
|
):
|
||||||
|
"""Initialize a PNG IHDR (Image Header) chunk.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
width (int): The width of the image in pixels.
|
||||||
|
height (int): The height of the image in pixels.
|
||||||
|
bit_depth (int): The number of bits per sample or per palette index (1, 2, 4, 8, or 16).
|
||||||
|
color_type (int): The color type of the image (0=Grayscale, 2=RGB, 3=Indexed,
|
||||||
|
4=Grayscale+Alpha, 6=RGBA).
|
||||||
|
crc (int | None, optional): CRC (Cyclic Redundancy Check) checksum for the chunk.
|
||||||
|
If None, it will be calculated automatically. Defaults
|
||||||
|
to None.
|
||||||
|
interlace (int, optional): The interlace method (0=no interlace, 1=Adam7 interlace).
|
||||||
|
Defaults to 0.
|
||||||
|
"""
|
||||||
|
self.width = width
|
||||||
|
self.height = height
|
||||||
|
self.bit_depth = bit_depth
|
||||||
|
self.color_type = color_type
|
||||||
|
self.crc = (
|
||||||
|
crc
|
||||||
|
or zlib.crc32(
|
||||||
|
b"IHDR"
|
||||||
|
+ width.to_bytes(4, "big")
|
||||||
|
+ height.to_bytes(4, "big")
|
||||||
|
+ bytes((bit_depth, color_type, 0, 0, interlace))
|
||||||
|
)
|
||||||
|
& 0xFFFFFFFF
|
||||||
|
)
|
||||||
|
self.interlace = interlace
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def from_packed(cls, packed: int, crc: int | None = None) -> Self:
|
||||||
|
"""Create a PNG instance from a packed integer representation.
|
||||||
|
|
||||||
|
Parse a packed integer containing width, height, bit depth, color type and interlace
|
||||||
|
bits and return a PNG instance. The CRC may be supplied; if omitted it will be
|
||||||
|
computed by the constructor.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
packed (int): Packed integer with fields laid out as described in the `packed` property.
|
||||||
|
crc (int | None): Optional CRC value for the IHDR chunk.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Self: PNG instance with parsed fields.
|
||||||
|
"""
|
||||||
|
interlace = packed & 0b1
|
||||||
|
color = (packed >> 1) & 0b111
|
||||||
|
depth = (packed >> 4) & 0b111
|
||||||
|
height = (packed >> 7) & 0xFFFF
|
||||||
|
width = (packed >> 23) & 0xFFFF
|
||||||
|
return cls(
|
||||||
|
width=width,
|
||||||
|
height=height,
|
||||||
|
bit_depth=depth,
|
||||||
|
color_type=color,
|
||||||
|
crc=crc,
|
||||||
|
interlace=interlace,
|
||||||
|
)
|
||||||
|
|
||||||
|
@cached_property
|
||||||
|
def packed(self) -> int:
|
||||||
|
"""Pack IHDR fields into a single integer.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
int: Packed integer encoding the IHDR fields.
|
||||||
|
|
||||||
|
Layout:
|
||||||
|
width : 16 bits
|
||||||
|
height : 16 bits
|
||||||
|
depth : 3 bits
|
||||||
|
color : 3 bits
|
||||||
|
interlace : 1 bit
|
||||||
|
"""
|
||||||
|
return (
|
||||||
|
((self.width & 0xFFFF) << 23)
|
||||||
|
| ((self.height & 0xFFFF) << 7)
|
||||||
|
| ((self.bit_depth & 0b111) << 4)
|
||||||
|
| ((self.color_type & 0b111) << 1)
|
||||||
|
| (self.interlace & 0b1)
|
||||||
|
)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def from_bytearray(cls, data: bytearray) -> Self:
|
||||||
|
"""Create a PNG image instance from a bytearray containing PNG file data.
|
||||||
|
|
||||||
|
Parses the PNG signature and IHDR (Image Header) chunk to extract image metadata
|
||||||
|
including dimensions, bit depth, color type, and interlacing information.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
data (bytearray): A bytearray containing the complete PNG file data, starting
|
||||||
|
with the PNG signature (8 bytes: \\x89PNG\\r\\n\\x1a\\n).
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Self | None: A new instance of the PNG image class with parsed IHDR metadata
|
||||||
|
(width, height, bit_depth, color_type, crc, interlace) if the data is
|
||||||
|
valid and contains a proper IHDR chunk. Returns None if the PNG signature
|
||||||
|
is invalid or the IHDR chunk type is not found.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
struct.error: If the data is too short or malformed for unpacking IHDR fields.
|
||||||
|
IndexError: If the bytearray is shorter than expected for the PNG structure.
|
||||||
|
"""
|
||||||
|
|
||||||
|
PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
|
||||||
|
if data[:8] != PNG_SIGNATURE:
|
||||||
|
raise ValueError("Invalid PNG signature")
|
||||||
|
offset = 8 # Read IHDR chunk
|
||||||
|
length = struct.unpack(">I", data[offset : offset + 4])[0]
|
||||||
|
offset += 4
|
||||||
|
chunk_type = data[offset : offset + 4]
|
||||||
|
offset += 4
|
||||||
|
if chunk_type != b"IHDR":
|
||||||
|
raise ValueError("IHDR chunk not found at expected position")
|
||||||
|
|
||||||
|
ihdr_data = data[offset : offset + length]
|
||||||
|
offset += length
|
||||||
|
|
||||||
|
crc = struct.unpack(">I", data[offset : offset + 4])[0]
|
||||||
|
|
||||||
|
# Unpack IHDR fields
|
||||||
|
width, height, depth, color, _, _, interlace = struct.unpack(
|
||||||
|
">IIBBBBB", ihdr_data
|
||||||
|
)
|
||||||
|
|
||||||
|
return cls(
|
||||||
|
width=width,
|
||||||
|
height=height,
|
||||||
|
bit_depth=depth,
|
||||||
|
color_type=color,
|
||||||
|
crc=crc,
|
||||||
|
interlace=interlace,
|
||||||
|
)
|
||||||
|
|
||||||
|
def __repr__(self) -> str:
|
||||||
|
return (
|
||||||
|
f"PNG(width={self.width}, height={self.height}, "
|
||||||
|
f"bit_depth={self.bit_depth}, color_type={self.color_type}, "
|
||||||
|
f"interlace={self.interlace}, crc=0x{self.crc:08x})"
|
||||||
|
)
|
||||||
|
|||||||
+22
-1
@@ -14,7 +14,28 @@ MAX_STORE_TIME = int(getenv("MAX_STORE_TIME", "259200")) # 3 days by default
|
|||||||
|
|
||||||
|
|
||||||
def cleanup_old_entries() -> None:
|
def cleanup_old_entries() -> None:
|
||||||
"""Clean up old entries in the database and file system."""
|
"""
|
||||||
|
Clean up old and incomplete entries from the database and file system.
|
||||||
|
|
||||||
|
This function performs two main cleanup operations:
|
||||||
|
|
||||||
|
1. Submission cleanup:
|
||||||
|
- Deletes submissions with "pending" or "running" status that have exceeded
|
||||||
|
the maximum allowed processing time (MAX_PENDING_TIME).
|
||||||
|
- Deletes completed submissions ("done" status) that have missing or buggy
|
||||||
|
results (results.json file not found).
|
||||||
|
|
||||||
|
2. Image cleanup:
|
||||||
|
- Deletes images older than MAX_STORE_TIME along with all associated
|
||||||
|
submissions and their result folders.
|
||||||
|
- Deletes orphaned images (with no submissions) older than MAX_PENDING_TIME
|
||||||
|
and removes their associated result folders from the file system.
|
||||||
|
|
||||||
|
All database changes are committed at the end of the operation.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
None
|
||||||
|
"""
|
||||||
now = time.time()
|
now = time.time()
|
||||||
for submission in Submission.query.all(): # type: ignore
|
for submission in Submission.query.all(): # type: ignore
|
||||||
if (
|
if (
|
||||||
|
|||||||
@@ -6,6 +6,3 @@ IMAGE_EXTENSIONS = [".png", ".jpg", ".jpeg", ".gif", ".bmp", ".webp", ".tiff"]
|
|||||||
WORKER_FILES = ["binwalk", "foremost", "steghide", "zsteg", "image_resize"]
|
WORKER_FILES = ["binwalk", "foremost", "steghide", "zsteg", "image_resize"]
|
||||||
|
|
||||||
RESULT_FOLDER = Path(__file__).parent.resolve() / "results"
|
RESULT_FOLDER = Path(__file__).parent.resolve() / "results"
|
||||||
|
|
||||||
MAX_WIDTH = 5 # Used for IHDR lookup table
|
|
||||||
MAX_HEIGHT = 5 # Used for IHDR lookup table
|
|
||||||
|
|||||||
+57
-51
@@ -1,16 +1,14 @@
|
|||||||
"""This module defines the database models for the Aperisolve application."""
|
"""This module defines the database models for the Aperi'Solve application."""
|
||||||
|
|
||||||
import zlib
|
|
||||||
import itertools
|
import itertools
|
||||||
|
|
||||||
from datetime import datetime, timezone
|
from datetime import datetime, timezone
|
||||||
from decimal import Decimal
|
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 import Flask
|
||||||
from flask_sqlalchemy import SQLAlchemy
|
from flask_sqlalchemy import SQLAlchemy
|
||||||
from sqlalchemy import Boolean, Column, DateTime, Float, Integer, String, BigInteger
|
from sqlalchemy import BigInteger, Boolean, Column, DateTime, Float, Integer, String
|
||||||
|
|
||||||
|
from aperisolve.analyzers.utils import PNG, get_resolutions, get_valid_depth_color_pairs
|
||||||
|
|
||||||
db: SQLAlchemy = SQLAlchemy()
|
db: SQLAlchemy = SQLAlchemy()
|
||||||
|
|
||||||
@@ -60,68 +58,76 @@ 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(
|
iid = Column(Integer, primary_key=True, autoincrement=True)
|
||||||
Integer, primary_key=True, autoincrement=True, unique=True
|
crc = Column(BigInteger, nullable=False)
|
||||||
)
|
packed = Column(BigInteger, nullable=False)
|
||||||
crc: Column[int] = Column(BigInteger)
|
|
||||||
packed: Column[int] = Column(Integer)
|
|
||||||
|
|
||||||
|
|
||||||
def create_crc_db() -> None:
|
def create_crc_db() -> None:
|
||||||
w = range(1, MAX_WIDTH + 1)
|
"""
|
||||||
h = range(1, MAX_HEIGHT + 1)
|
Creates a database of common PNG IHDR (Image Header) entries by generating
|
||||||
|
combinations of standard PNG parameters such as resolutions, bit depths,
|
||||||
|
color types, and interlace methods. The function checks for existing entries
|
||||||
|
in the database to avoid duplicates and inserts new entries in batches of
|
||||||
|
5000, committing the changes to the database periodically.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
None
|
||||||
|
"""
|
||||||
# Standard PNG IHDR parameters
|
# Standard PNG IHDR parameters
|
||||||
bit_depth = [1, 2, 4, 8, 16]
|
resolutions = get_resolutions()
|
||||||
color_type = [0, 2, 3, 4, 6]
|
bit_color_pairs = get_valid_depth_color_pairs()
|
||||||
compression_method = [0] # This is always 0 according to the PNG spec
|
interlace_methods = [0, 1]
|
||||||
filter_method = [0] # This is always 0 according to the PNG spec
|
|
||||||
interlace_method = [0, 1]
|
|
||||||
|
|
||||||
# Create the generator
|
combinations = itertools.product(
|
||||||
all_combinations = itertools.product(
|
resolutions,
|
||||||
w, h, bit_depth, color_type, compression_method, filter_method, interlace_method
|
bit_color_pairs,
|
||||||
|
interlace_methods,
|
||||||
)
|
)
|
||||||
|
|
||||||
count = 0
|
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
|
with db.session.no_autoflush: # pylint: disable=no-member
|
||||||
crc = zlib.crc32(b"IHDR" + ihdr_data) & 0xFFFFFFFF
|
for (w, h), (bd, ct), inter in combinations:
|
||||||
|
p = PNG(width=w, height=h, bit_depth=bd, color_type=ct, interlace=inter)
|
||||||
|
|
||||||
# Compute packed
|
exists = db.session.execute( # pylint: disable=no-member
|
||||||
packed = pack_ihdr(w_val, h_val, bd, ct, inter)
|
db.select(IHDR.iid).where(IHDR.packed == p.packed)
|
||||||
|
).first()
|
||||||
|
|
||||||
# Create and add IHDR entry to database
|
if exists: # Skip if already present
|
||||||
ihdr_entry = IHDR(crc=crc, packed=packed)
|
continue
|
||||||
db.session.add(ihdr_entry)
|
|
||||||
|
|
||||||
if count % 10000 == 0:
|
db.session.add(IHDR(packed=p.packed, crc=p.crc)) # pylint: disable=no-member
|
||||||
print(
|
count += 1
|
||||||
f"Inserted {count}/{max_count} ({int((count*100)/max_count)}%) IHDR entries..."
|
|
||||||
)
|
|
||||||
db.session.commit()
|
|
||||||
|
|
||||||
# Commit all entries
|
if count % 5000 == 0:
|
||||||
db.session.commit()
|
print(f"Inserted {count} common IHDR entries...")
|
||||||
|
db.session.commit()
|
||||||
|
|
||||||
|
db.session.commit() # pylint: disable=no-member
|
||||||
|
print(f"Precomputed {count} common IHDR entries.")
|
||||||
|
|
||||||
|
|
||||||
def init_db(app: Flask) -> None:
|
def init_db(app: Flask) -> None:
|
||||||
"""Initialize the database with the given Flask app."""
|
"""
|
||||||
|
Initialize the database by creating tables and populating lookup data.
|
||||||
|
|
||||||
|
This function creates the database schema if it doesn't already exist,
|
||||||
|
and populates the CRC lookup table for PNG IHDR validation.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
app (Flask): The Flask application instance used to establish the application context
|
||||||
|
for database operations.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
None
|
||||||
|
|
||||||
|
Side Effects:
|
||||||
|
- Creates all database tables if the "image" table doesn't exist
|
||||||
|
- Populates the IHDR CRC lookup table with initial data
|
||||||
|
- Prints status messages to console during initialization
|
||||||
|
"""
|
||||||
with app.app_context():
|
with app.app_context():
|
||||||
if not db.engine.dialect.has_table(db.engine.connect(), "image"):
|
if not db.engine.dialect.has_table(db.engine.connect(), "image"):
|
||||||
print("Creating database...")
|
print("Creating database...")
|
||||||
|
|||||||
+29
-5
@@ -8,19 +8,43 @@ from .analyzers.binwalk import analyze_binwalk
|
|||||||
from .analyzers.decomposer import analyze_decomposer
|
from .analyzers.decomposer import analyze_decomposer
|
||||||
from .analyzers.exiftool import analyze_exiftool
|
from .analyzers.exiftool import analyze_exiftool
|
||||||
from .analyzers.foremost import analyze_foremost
|
from .analyzers.foremost import analyze_foremost
|
||||||
|
from .analyzers.image_resize import analyze_image_resize
|
||||||
from .analyzers.outguess import analyze_outguess
|
from .analyzers.outguess import analyze_outguess
|
||||||
|
from .analyzers.pngcheck import analyze_pngcheck
|
||||||
from .analyzers.steghide import analyze_steghide
|
from .analyzers.steghide import analyze_steghide
|
||||||
from .analyzers.strings import analyze_strings
|
from .analyzers.strings import analyze_strings
|
||||||
from .analyzers.zsteg import analyze_zsteg
|
from .analyzers.zsteg import analyze_zsteg
|
||||||
from .analyzers.pngcheck import analyze_pngcheck
|
|
||||||
from .analyzers.image_resize import analyze_image_resize
|
|
||||||
from .app import app, db
|
from .app import app, db
|
||||||
from .config import RESULT_FOLDER
|
from .config import RESULT_FOLDER
|
||||||
from .models import Image, Submission
|
from .models import Image, Submission
|
||||||
|
|
||||||
|
|
||||||
def analyze_image(submission_hash: str) -> None:
|
def analyze_image(submission_hash: str) -> None:
|
||||||
"""Analyze an image submission in separate threads."""
|
"""
|
||||||
|
Analyze an image submission by running multiple analysis tools concurrently.
|
||||||
|
|
||||||
|
This function retrieves a submission and its associated image from the database,
|
||||||
|
then executes various image analysis tools in parallel threads. The analysis tools
|
||||||
|
include binwalk, decomposer, exiftool, foremost, strings, pngcheck, steghide,
|
||||||
|
zsteg, and image resize. If deep analysis is enabled, additional tools like
|
||||||
|
outguess are included.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
submission_hash (str): The unique hash identifier of the submission to analyze.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
None
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
Implicitly catches and handles exceptions during analysis, setting submission
|
||||||
|
status to "error" if an exception occurs.
|
||||||
|
|
||||||
|
Side Effects:
|
||||||
|
- Queries the database for the submission and associated image
|
||||||
|
- Updates submission status in the database ("running" -> "completed" or "error")
|
||||||
|
- Creates result directories and generates analysis output files
|
||||||
|
- Modifies the database session and commits changes
|
||||||
|
"""
|
||||||
with app.app_context():
|
with app.app_context():
|
||||||
submission: Submission = Submission.query.get(submission_hash) # type: ignore
|
submission: Submission = Submission.query.get(submission_hash) # type: ignore
|
||||||
image = Image.query.get_or_404(submission.image_hash) # type: ignore
|
image = Image.query.get_or_404(submission.image_hash) # type: ignore
|
||||||
@@ -29,7 +53,7 @@ def analyze_image(submission_hash: str) -> None:
|
|||||||
return
|
return
|
||||||
|
|
||||||
submission.status = "running" # type: ignore
|
submission.status = "running" # type: ignore
|
||||||
db.session.commit()
|
db.session.commit() # pylint: disable=no-member
|
||||||
|
|
||||||
try:
|
try:
|
||||||
img_path: Path = Path(image.file)
|
img_path: Path = Path(image.file)
|
||||||
@@ -77,4 +101,4 @@ def analyze_image(submission_hash: str) -> None:
|
|||||||
except Exception:
|
except Exception:
|
||||||
submission.status = "error"
|
submission.status = "error"
|
||||||
finally:
|
finally:
|
||||||
db.session.commit()
|
db.session.commit() # pylint: disable=no-member
|
||||||
|
|||||||
Reference in New Issue
Block a user