Image resize and linting
This commit is contained in:
Zeecka
@@ -1,17 +1,15 @@
|
||||
"""Image size bruteforcer for common PNG sizes."""
|
||||
|
||||
import zlib
|
||||
import sqlite3
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
from .utils import update_data, unpack_ihdr
|
||||
|
||||
from ..app import app
|
||||
from ..models import IHDR
|
||||
from .utils import PNG, update_data
|
||||
|
||||
|
||||
def write_recovered_png(
|
||||
png_bytes: bytearray,
|
||||
ihdr_offset: int,
|
||||
width: int,
|
||||
height: int,
|
||||
output_path: Path,
|
||||
@@ -19,40 +17,102 @@ 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")
|
||||
|
||||
# Splicing: [Start...IHDR+4] + [W] + [H] + [IHDR+12...End]
|
||||
full_png_data = (
|
||||
png_bytes[: ihdr_offset + 4]
|
||||
png_bytes[:16] # Splicing: [Start...IHDR+4] + [W] + [H] + [IHDR+12...End]
|
||||
+ width_bytes
|
||||
+ height_bytes
|
||||
+ png_bytes[ihdr_offset + 12 :]
|
||||
+ png_bytes[24:]
|
||||
)
|
||||
|
||||
with output_path.open("wb") as out_f:
|
||||
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():
|
||||
crc_b = int.from_bytes(crc_bytes, byteorder="big")
|
||||
imgs = IHDR.query.filter_by(crc=crc_b).all()
|
||||
imgs = IHDR.query.filter_by(crc=png.crc).all()
|
||||
results = []
|
||||
for img in imgs:
|
||||
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
|
||||
|
||||
|
||||
def analyze_image_resize(input_img: Path, output_dir: Path) -> None:
|
||||
"""
|
||||
Analyze an image submission using python-based resize bruteforce.
|
||||
Strategy: Check common sizes first, then fall back to DB lookup.
|
||||
Analyze a PNG image and recover resized versions based on CRC matching.
|
||||
|
||||
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
|
||||
@@ -64,48 +124,48 @@ def analyze_image_resize(input_img: Path, output_dir: Path) -> None:
|
||||
|
||||
try:
|
||||
with input_img.open("rb") as image:
|
||||
f = image.read()
|
||||
img_bytes = bytearray(image.read())
|
||||
png = PNG.from_bytearray(img_bytes)
|
||||
|
||||
b = bytearray(f)
|
||||
|
||||
# Find IHDR start
|
||||
ihdr = b.find(b"\x49\x48\x44\x52")
|
||||
if ihdr == -1:
|
||||
logs.append("Failure: PNG header (IHDR) not found.")
|
||||
if png is None:
|
||||
logs.append(
|
||||
"Failure: IHDR chunk is not the first chunk, or PNG has invalid structure."
|
||||
)
|
||||
update_data(
|
||||
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
|
||||
|
||||
# Calculate Chunk Length (The 4 bytes BEFORE IHDR tag)
|
||||
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()}")
|
||||
logs.append(f"Target CRC found: 0x{png.crc:08x}")
|
||||
|
||||
saved_img_urls = []
|
||||
|
||||
db_matches = lookup_crc(target_crc)
|
||||
if db_matches:
|
||||
db_matches = lookup_crc(png)
|
||||
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
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Iterate through ALL candidates found
|
||||
for elt in db_matches:
|
||||
w, h, _, _, _ = elt # TODO, use depth/color later if needed
|
||||
img_name = f"recovered_{w}x{h}.png"
|
||||
for im_match in db_matches:
|
||||
img_name = f"recovered_{im_match.width}x{im_match.height}.png"
|
||||
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}")
|
||||
|
||||
# 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 = {
|
||||
|
||||
+206
-17
@@ -2,13 +2,14 @@
|
||||
|
||||
import fcntl
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import struct
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
import itertools
|
||||
import zlib
|
||||
import sqlite3
|
||||
from functools import cached_property
|
||||
from pathlib import Path
|
||||
from typing import Any, Self
|
||||
|
||||
_thread_lock = threading.Lock()
|
||||
|
||||
@@ -54,9 +55,147 @@ def update_data(
|
||||
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 = (
|
||||
list(range(16, 257, 16)) # 16 → 256
|
||||
+ list(range(320, 1025, 32)) # 320 → 1024
|
||||
+ list(range(1280, 2561, 64)) # 1280 → 2560
|
||||
+ list(range(3000, 4097, 128)) # 3K–4K
|
||||
+ list(range(5120, 8193, 256)) # 5K–8K
|
||||
+ [10000] # upper bound
|
||||
)
|
||||
ASPECT_RATIOS = [
|
||||
# Screens / digital
|
||||
(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 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:
|
||||
"""
|
||||
Pack IHDR fields into a single integer.
|
||||
A class representing a PNG image header (IHDR chunk).
|
||||
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.
|
||||
"""
|
||||
|
||||
width: int
|
||||
height: int
|
||||
bit_depth: int
|
||||
color_type: int
|
||||
crc: int
|
||||
interlace: int = 0
|
||||
|
||||
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
|
||||
@@ -65,19 +204,69 @@ def pack_ihdr(width: int, height: int, depth: int, color: int, interlace: int) -
|
||||
color : 3 bits
|
||||
interlace : 1 bit
|
||||
"""
|
||||
return (width << 23) | (height << 7) | (depth << 4) | (color << 1) | interlace
|
||||
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.
|
||||
|
||||
def unpack_ihdr(packed: int) -> tuple[int, int, int, int, int]:
|
||||
"""
|
||||
Unpack packed IHDR integer into components:
|
||||
(width, height, depth, color, interlace)
|
||||
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.
|
||||
"""
|
||||
|
||||
interlace = packed & 0b1
|
||||
color = (packed >> 1) & 0b111
|
||||
depth = (packed >> 4) & 0b111
|
||||
height = (packed >> 7) & 0xFFFF
|
||||
width = (packed >> 23) & 0xFFFF
|
||||
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")
|
||||
|
||||
return width, height, depth, color, interlace
|
||||
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:
|
||||
"""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()
|
||||
for submission in Submission.query.all(): # type: ignore
|
||||
if (
|
||||
|
||||
@@ -6,6 +6,3 @@ 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 = 5 # Used for IHDR lookup table
|
||||
MAX_HEIGHT = 5 # Used for IHDR lookup table
|
||||
|
||||
+58
-52
@@ -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
|
||||
|
||||
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, 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()
|
||||
|
||||
@@ -60,68 +58,76 @@ 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)
|
||||
packed: Column[int] = Column(Integer)
|
||||
iid = Column(Integer, primary_key=True, autoincrement=True)
|
||||
crc = Column(BigInteger, nullable=False)
|
||||
packed = Column(BigInteger, nullable=False)
|
||||
|
||||
|
||||
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
|
||||
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]
|
||||
resolutions = get_resolutions()
|
||||
bit_color_pairs = get_valid_depth_color_pairs()
|
||||
interlace_methods = [0, 1]
|
||||
|
||||
# Create the generator
|
||||
all_combinations = itertools.product(
|
||||
w, h, bit_depth, color_type, compression_method, filter_method, interlace_method
|
||||
combinations = itertools.product(
|
||||
resolutions,
|
||||
bit_color_pairs,
|
||||
interlace_methods,
|
||||
)
|
||||
|
||||
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:
|
||||
|
||||
with db.session.no_autoflush: # pylint: disable=no-member
|
||||
for (w, h), (bd, ct), inter in combinations:
|
||||
p = PNG(width=w, height=h, bit_depth=bd, color_type=ct, interlace=inter)
|
||||
|
||||
exists = db.session.execute( # pylint: disable=no-member
|
||||
db.select(IHDR.iid).where(IHDR.packed == p.packed)
|
||||
).first()
|
||||
|
||||
if exists: # Skip if already present
|
||||
continue
|
||||
|
||||
db.session.add(IHDR(packed=p.packed, crc=p.crc)) # pylint: disable=no-member
|
||||
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..."
|
||||
)
|
||||
if count % 5000 == 0:
|
||||
print(f"Inserted {count} common IHDR entries...")
|
||||
db.session.commit()
|
||||
|
||||
# Commit all entries
|
||||
db.session.commit()
|
||||
db.session.commit() # pylint: disable=no-member
|
||||
print(f"Precomputed {count} common IHDR entries.")
|
||||
|
||||
|
||||
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():
|
||||
if not db.engine.dialect.has_table(db.engine.connect(), "image"):
|
||||
print("Creating database...")
|
||||
|
||||
+29
-5
@@ -8,19 +8,43 @@ from .analyzers.binwalk import analyze_binwalk
|
||||
from .analyzers.decomposer import analyze_decomposer
|
||||
from .analyzers.exiftool import analyze_exiftool
|
||||
from .analyzers.foremost import analyze_foremost
|
||||
from .analyzers.image_resize import analyze_image_resize
|
||||
from .analyzers.outguess import analyze_outguess
|
||||
from .analyzers.pngcheck import analyze_pngcheck
|
||||
from .analyzers.steghide import analyze_steghide
|
||||
from .analyzers.strings import analyze_strings
|
||||
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 .config import RESULT_FOLDER
|
||||
from .models import Image, Submission
|
||||
|
||||
|
||||
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():
|
||||
submission: Submission = Submission.query.get(submission_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
|
||||
|
||||
submission.status = "running" # type: ignore
|
||||
db.session.commit()
|
||||
db.session.commit() # pylint: disable=no-member
|
||||
|
||||
try:
|
||||
img_path: Path = Path(image.file)
|
||||
@@ -77,4 +101,4 @@ def analyze_image(submission_hash: str) -> None:
|
||||
except Exception:
|
||||
submission.status = "error"
|
||||
finally:
|
||||
db.session.commit()
|
||||
db.session.commit() # pylint: disable=no-member
|
||||
|
||||
Reference in New Issue
Block a user