image_file_budget.py

"""Bounded image container inspection; never decompress pixels or retain image payloads."""
import os
import time

MAX_EDGE = 16_000
MAX_PIXELS = 40_000_000
MAX_FRAMES = 240
MAX_ANIMATION_PIXELS = 400_000_000
MAX_FILE_BYTES = 50 * 1024 * 1024
MAX_STEPS = 500_000
SCAN_SECONDS = 3
JPEG_HEADER_BYTES = 256 * 1024

class ImageBudgetError(ValueError):
    pass

def _bad():
    raise ImageBudgetError('invalid or over-budget image container')

def _budget(width, height, frames=1):
    if not 0 < width <= MAX_EDGE or not 0 < height <= MAX_EDGE or width * height > MAX_PIXELS:
        _bad()
    if not 0 < frames <= MAX_FRAMES or width * height * frames > MAX_ANIMATION_PIXELS:
        _bad()

class _Reader:
    def __init__(self, file):
        self.file = file
        self.size = os.fstat(file.fileno()).st_size
        self.deadline = time.monotonic() + SCAN_SECONDS
        self.steps = 0
        if not 0 < self.size <= MAX_FILE_BYTES: _bad()
    def check(self, length):
        self.steps += 1
        if self.steps > MAX_STEPS or time.monotonic() > self.deadline or length < 0 or self.file.tell() + length > self.size:
            _bad()
    def read(self, length):
        self.check(length)
        if length > 32: _bad()
        result = self.file.read(length)
        if len(result) != length: _bad()
        return result
    def skip(self, length):
        self.check(length)
        self.file.seek(length, 1)
    def pos(self): return self.file.tell()

def _png(r):
    if r.read(8) != b'\x89PNG\r\n\x1a\n': _bad()
    if r.read(8) != b'\x00\x00\x00\rIHDR': _bad()
    header = r.read(13)
    width, height = int.from_bytes(header[:4], 'big'), int.from_bytes(header[4:8], 'big')
    _budget(width, height)
    r.skip(4)
    declared, frames, data, first_control_before_data = None, 0, False, False
    while r.pos() < r.size:
        head = r.read(8); length, kind = int.from_bytes(head[:4], 'big'), head[4:]
        if kind == b'acTL':
            if declared is not None or data or length != 8: _bad()
            declared = int.from_bytes(r.read(8)[:4], 'big'); _budget(width, height, declared)
        elif kind == b'fcTL':
            if declared is None or length != 26: _bad()
            control = r.read(26)
            w, h, x, y = [int.from_bytes(control[n:n+4], 'big') for n in (4,8,12,16)]
            if not w or not h or x + w > width or y + h > height: _bad()
            if frames == 0: first_control_before_data = not data
            frames += 1; _budget(width, height, frames)
        elif kind == b'IDAT':
            data = True; r.skip(length)
        elif kind == b'IEND':
            if length != 0 or not data or (declared is not None and frames != declared): _bad()
            r.skip(4)
            if r.pos() != r.size: _bad()
            total = frames + (0 if first_control_before_data else 1) if declared is not None else 1
            _budget(width, height, total)
            return width, height, total
        else:
            r.skip(length)
        r.skip(4)
    _bad()

def _gif(r):
    if r.read(6) not in (b'GIF87a', b'GIF89a'): _bad()
    head = r.read(7); width, height = int.from_bytes(head[:2],'little'), int.from_bytes(head[2:4],'little')
    _budget(width,height)
    if head[4] & 128: r.skip(3 * (1 << ((head[4] & 7) + 1)))
    def blocks():
        while True:
            size = r.read(1)[0]
            if not size: return
            r.skip(size)
    frames = 0
    while r.pos() < r.size:
        kind = r.read(1)[0]
        if kind == 0x21:
            r.read(1); blocks()
        elif kind == 0x2c:
            image = r.read(9)
            x,y,w,h = [int.from_bytes(image[n:n+2],'little') for n in (0,2,4,6)]
            if not w or not h or x+w > width or y+h > height: _bad()
            frames += 1; _budget(width,height,frames)
            if image[8] & 128: r.skip(3 * (1 << ((image[8] & 7) + 1)))
            if not 2 <= r.read(1)[0] <= 8: _bad()
            blocks()
        elif kind == 0x3b:
            if not frames or r.pos() != r.size: _bad()
            return width,height,frames
        else: _bad()
    _bad()

def _jpeg(r):
    if r.read(2) != b'\xff\xd8': _bad()
    while r.pos() < min(r.size, JPEG_HEADER_BYTES):
        if r.read(1) != b'\xff': _bad()
        marker = r.read(1)[0]
        while marker == 0xff: marker = r.read(1)[0]
        if marker in (0,0xd8,0xd9,0xda) or 0xd0 <= marker <= 0xd7: _bad()
        length = int.from_bytes(r.read(2),'big')
        if length < 2 or r.pos() + length - 2 > min(r.size, JPEG_HEADER_BYTES): _bad()
        if 0xc0 <= marker <= 0xcf and marker not in (0xc4,0xc8,0xcc):
            if length < 8: _bad()
            head = r.read(6); height, width = int.from_bytes(head[1:3],'big'), int.from_bytes(head[3:5],'big')
            _budget(width,height)
            return width,height,1
        r.skip(length-2)
    _bad()

def _webp(r):
    head = r.read(12)
    if head[:4] != b'RIFF' or head[8:] != b'WEBP' or int.from_bytes(head[4:8],'little')+8 != r.size: _bad()
    canvas, frames, animated, animation_header, static = None, 0, False, False, False
    def codec(kind,length):
        if kind == b'VP8 ':
            if length < 10: _bad()
            h = r.read(10)
            if h[3:6] != b'\x9d\x01\x2a' or h[0] & 1: _bad()
            result = int.from_bytes(h[6:8],'little') & 0x3fff, int.from_bytes(h[8:10],'little') & 0x3fff
            r.skip(length-10)
        elif kind == b'VP8L':
            if length < 5: _bad()
            h = r.read(5)
            if h[0] != 0x2f: _bad()
            bits = int.from_bytes(h[1:],'little')
            if bits >> 29: _bad()
            result = (bits & 0x3fff)+1, ((bits >> 14) & 0x3fff)+1
            r.skip(length-5)
        else: _bad()
        _budget(*result)
        return result
    while r.pos() < r.size:
        head = r.read(8); kind, length = head[:4], int.from_bytes(head[4:],'little')
        r.check(length+(length & 1)); end = r.pos()+length
        if kind == b'VP8X':
            if canvas is not None or length != 10: _bad()
            h=r.read(10); canvas=(int.from_bytes(h[4:7],'little')+1,int.from_bytes(h[7:10],'little')+1)
            _budget(*canvas); animated=bool(h[0] & 2)
        elif kind == b'ANIM':
            if not animated or length != 6 or animation_header: _bad()
            animation_header=True; r.skip(length)
        elif kind == b'ANMF':
            if not animated or not animation_header or canvas is None or length < 16: _bad()
            h=r.read(16); x,y=int.from_bytes(h[:3],'little')*2,int.from_bytes(h[3:6],'little')*2
            w,hh=int.from_bytes(h[6:9],'little')+1,int.from_bytes(h[9:12],'little')+1
            if x+w > canvas[0] or y+hh > canvas[1]: _bad()
            frames+=1; _budget(*canvas,frames)
            found=False
            while r.pos()<end:
                chunk=r.read(8); sub,size=chunk[:4],int.from_bytes(chunk[4:],'little')
                if r.pos()+size+(size & 1)>end: _bad()
                if sub in (b'VP8 ',b'VP8L'):
                    if found or codec(sub,size)!=(w,hh): _bad()
                    found=True
                else: r.skip(size)
                if size & 1: r.skip(1)
            if not found: _bad()
        elif kind in (b'VP8 ',b'VP8L'):
            if animated or static: _bad()
            size=codec(kind,length)
            if canvas is not None and size!=canvas: _bad()
            canvas=size;static=True
        else: r.skip(length)
        if r.pos()!=end: _bad()
        if length & 1: r.skip(1)
    if canvas is None or (animated and not frames) or (not animated and not static): _bad()
    return canvas[0], canvas[1], frames if animated else 1

def inspect_image_file(path, ext):
    parser={'png':_png,'gif':_gif,'jpg':_jpeg,'jpeg':_jpeg,'webp':_webp}.get(ext)
    if parser is None: _bad()
    with open(path,'rb') as file:
        width,height,frames=parser(_Reader(file))
    return {'width':width,'height':height,'frames':frames}