Spaces:
Running
Running
| """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} | |