"""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: _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}