105 lines
4.3 KiB
Python
105 lines
4.3 KiB
Python
"""MPF fixture: Pillow MPO, two generated 8x8 red/blue frames.
|
|
|
|
Generate with Image.new('RGB', (8, 8), color) for red, blue;
|
|
first.save('multipicture.jpg', format='MPO', save_all=True,
|
|
append_images=remaining). No private photographs are used.
|
|
"""
|
|
import base64
|
|
from pathlib import Path
|
|
import struct
|
|
|
|
import pytest
|
|
|
|
from api import helpers
|
|
|
|
|
|
RAW = (Path(__file__).parent / 'fixtures' / 'multipicture.jpg').read_bytes()
|
|
BASE = RAW.index(b'MPF\0') + 4
|
|
IFD = BASE + int.from_bytes(RAW[BASE + 4:BASE + 8], 'little')
|
|
TAGS = {int.from_bytes(RAW[p:p + 2], 'little'): p
|
|
for p in range(IFD + 2, IFD + 2 + 12 * int.from_bytes(RAW[IFD:IFD + 2], 'little'), 12)}
|
|
ENTRIES = BASE + int.from_bytes(RAW[TAGS[0xB002] + 8:TAGS[0xB002] + 12], 'little')
|
|
|
|
|
|
def uri(raw):
|
|
return 'data:image/jpeg;base64,' + base64.b64encode(raw).decode('ascii')
|
|
|
|
|
|
def session(value):
|
|
return {'messages': [{'role': 'user', 'content': [
|
|
{'type': 'image_url', 'image_url': {'url': value}}]}]}
|
|
|
|
|
|
@pytest.mark.parametrize('byte_order', ['little', 'big'])
|
|
def test_mpf_preserves_native_image_without_text_redaction(monkeypatch, byte_order):
|
|
raw = bytearray(RAW)
|
|
if byte_order == 'big':
|
|
raw[BASE:BASE + 4] = b'MM\x00\x2a'
|
|
struct.pack_into('>I', raw, BASE + 4, IFD - BASE)
|
|
struct.pack_into('>H', raw, IFD, len(TAGS))
|
|
for tag, pos in TAGS.items():
|
|
kind, count, value = struct.unpack_from('<HII', RAW, pos + 2)
|
|
struct.pack_into('>HHI', raw, pos, tag, kind, count)
|
|
if tag != 0xB000: # Version is four raw ASCII bytes, not an integer.
|
|
struct.pack_into('>I', raw, pos + 8, value)
|
|
for pos in range(ENTRIES, ENTRIES + 32, 16):
|
|
struct.pack_into('>IIIHH', raw, pos, *struct.unpack_from('<IIIHH', RAW, pos))
|
|
value = uri(raw)
|
|
calls = []
|
|
monkeypatch.setattr(helpers, '_redact_text', lambda text, **kw: calls.append(text) or text)
|
|
assert helpers._is_native_raster_data_uri(value)
|
|
assert helpers.redact_session_data(session(value)) == session(value)
|
|
assert value not in calls
|
|
|
|
|
|
@pytest.mark.parametrize('attack', ['trailing_secret', 'gap', 'overlap', 'size',
|
|
'count', 'table_offset', 'truncated', 'bad_frame',
|
|
'first_offset', 'undeclared_image'])
|
|
def test_mpf_rejects_unaccounted_or_invalid_bytes(attack):
|
|
raw = bytearray(RAW)
|
|
if attack == 'trailing_secret':
|
|
raw += b'API_KEY=private-trailing-secret'
|
|
elif attack == 'undeclared_image':
|
|
raw += RAW
|
|
elif attack == 'truncated':
|
|
del raw[-1:]
|
|
elif attack == 'count':
|
|
struct.pack_into('<I', raw, TAGS[0xB001] + 8, 999999)
|
|
elif attack == 'table_offset':
|
|
struct.pack_into('<I', raw, TAGS[0xB002] + 8, len(raw))
|
|
elif attack == 'first_offset':
|
|
struct.pack_into('<I', raw, ENTRIES + 8, 1)
|
|
elif attack == 'bad_frame':
|
|
raw[-2:] = b'xx'
|
|
else:
|
|
pos = ENTRIES + 16 + (4 if attack == 'size' else 8)
|
|
value = int.from_bytes(raw[pos:pos + 4], 'little')
|
|
struct.pack_into('<I', raw, pos, value + (-1 if attack == 'overlap' else 1))
|
|
assert not helpers._is_native_raster_data_uri(uri(raw))
|
|
|
|
|
|
def test_mpf_appended_base64_secret_reaches_real_redactor():
|
|
# Pad inside a declared JPEG COM segment, keeping all MP offsets valid,
|
|
# so the appended decoded credential aligns to a base64 quantum.
|
|
raw = bytearray(RAW)
|
|
padding = b'\xff\xfe\x00\x04xx' if len(raw) % 3 == 0 else (
|
|
b'\xff\xfe\x00\x02' if len(raw) % 3 == 2 else b'\xff\xfe\x00\x03x')
|
|
raw[-2:-2] = padding
|
|
last_size = ENTRIES + 16 + 4
|
|
struct.pack_into('<I', raw, last_size, int.from_bytes(RAW[last_size:last_size + 4], 'little') + len(padding))
|
|
assert len(raw) % 3 == 0
|
|
secret = 'AKIA' + 'Z' * 16
|
|
value = uri(raw + base64.b64decode(secret))
|
|
assert value.endswith(secret)
|
|
assert not helpers._is_native_raster_data_uri(value)
|
|
result = helpers.redact_session_data(session(value))
|
|
assert secret not in result['messages'][0]['content'][0]['image_url']['url']
|
|
|
|
|
|
def test_mpf_outside_authoritative_image_boundary_still_redacts(monkeypatch):
|
|
value = uri(RAW)
|
|
calls = []
|
|
monkeypatch.setattr(helpers, '_redact_text', lambda text, **kw: calls.append(text) or text)
|
|
helpers.redact_session_data({'tool_calls': [{'args': session(value)}]})
|
|
assert value in calls
|