118 lines
6.2 KiB
Python
118 lines
6.2 KiB
Python
import io
|
|
import unittest
|
|
import zlib
|
|
from unittest.mock import patch
|
|
from pypdf import PdfWriter, PdfReader
|
|
from pypdf.generic import NameObject, DictionaryObject, DecodedStreamObject, EncodedStreamObject, TextStringObject
|
|
from engine import inspect_content
|
|
from pdf_documents import extract_pdf
|
|
|
|
|
|
def policy():
|
|
return {"dlp": {"maxBytes": 1048576, "rules": [{"id": "secret", "kind": "literal", "value": "Secret Roadmap"}], "semantic": {"enabled": False}}}
|
|
|
|
|
|
def pdf(value=b'BT /F1 24 Tf 20 150 Td (Public information) Tj ET', mutate=None):
|
|
writer = PdfWriter()
|
|
page = writer.add_blank_page(400, 220)
|
|
font = DictionaryObject({NameObject('/Type'): NameObject('/Font'), NameObject('/Subtype'): NameObject('/Type1'), NameObject('/BaseFont'): NameObject('/Helvetica')})
|
|
page[NameObject('/Resources')] = DictionaryObject({NameObject('/Font'): DictionaryObject({NameObject('/F1'): writer._add_object(font)})})
|
|
stream = DecodedStreamObject(); stream.set_data(value)
|
|
page[NameObject('/Contents')] = writer._add_object(stream)
|
|
if mutate: mutate(writer, page)
|
|
output = io.BytesIO(); writer.write(output)
|
|
return output.getvalue()
|
|
|
|
|
|
class PdfTests(unittest.TestCase):
|
|
def test_public_render_and_ocr_all_pages(self):
|
|
result = inspect_content(pdf(), policy())
|
|
self.assertEqual(result['decision'], 'allow')
|
|
self.assertTrue(result['coverage']['complete'])
|
|
self.assertEqual(result['coverage']['mediaItems'], 1)
|
|
|
|
def test_encoded_secret_decoded_text(self):
|
|
data = pdf(b'BT /F1 24 Tf 20 150 Td <53656372657420526F61646D6170> Tj ET')
|
|
self.assertNotIn(b'Secret Roadmap', data)
|
|
self.assertEqual(inspect_content(data, policy())['ruleId'], 'secret')
|
|
|
|
def test_hidden_metadata_and_unreferenced_strings(self):
|
|
for mutation in [lambda w, _: w.add_metadata({'/Title': 'Secret Roadmap'}), lambda w, _: w._add_object(TextStringObject('Secret Roadmap'))]:
|
|
self.assertEqual(inspect_content(pdf(mutate=mutation), policy())['ruleId'], 'secret')
|
|
|
|
def test_render_only_vector_secret_has_no_extractable_text(self):
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
mask = Image.new('L', (500, 80), 255)
|
|
ImageDraw.Draw(mask).text((10, 8), 'Secret Roadmap', font=ImageFont.load_default(size=42), fill=0)
|
|
operations = []
|
|
for y in range(mask.height):
|
|
x = 0
|
|
while x < mask.width:
|
|
if mask.getpixel((x, y)) < 128:
|
|
start = x
|
|
while x < mask.width and mask.getpixel((x, y)) < 128: x += 1
|
|
operations.append(f'{start + 15} {180-y} {x-start} 1 re f')
|
|
else: x += 1
|
|
data = pdf(('\n'.join(operations)).encode())
|
|
self.assertFalse(PdfReader(io.BytesIO(data)).pages[0].extract_text())
|
|
self.assertNotIn(b'Secret Roadmap', data)
|
|
self.assertEqual(inspect_content(data, policy())['ruleId'], 'secret')
|
|
|
|
def test_embedded_raster_secret(self):
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
image = Image.new('RGB', (700, 200), 'white')
|
|
ImageDraw.Draw(image).text((20, 40), 'Secret Roadmap', font=ImageFont.load_default(size=48), fill='black')
|
|
output = io.BytesIO(); image.save(output, format='PDF')
|
|
self.assertEqual(inspect_content(output.getvalue(), policy())['ruleId'], 'secret')
|
|
|
|
def test_unsupported_active_encrypted_or_hidden_stream(self):
|
|
def action(writer, _): writer._root_object[NameObject('/OpenAction')] = DictionaryObject({NameObject('/S'): NameObject('/JavaScript')})
|
|
def attachment(writer, _): writer.add_attachment('private.txt', b'opaque')
|
|
def encryption(writer, _): writer.encrypt('private')
|
|
def hidden_stream(writer, _):
|
|
stream = DecodedStreamObject(); stream.set_data(b'opaque'); writer._add_object(stream)
|
|
for mutation in [action, attachment, encryption, hidden_stream]:
|
|
result = inspect_content(pdf(mutate=mutation), policy())
|
|
self.assertEqual(result['decision'], 'block')
|
|
self.assertFalse(result['coverage']['complete'])
|
|
|
|
def test_decoder_and_render_budgets_fail_before_ocr(self):
|
|
def bomb(writer, page):
|
|
stream = EncodedStreamObject(); stream._data = zlib.compress(b'x' * (21 * 1024 * 1024))
|
|
stream[NameObject('/Filter')] = NameObject('/FlateDecode')
|
|
reference = page.raw_get('/Contents')
|
|
stream.indirect_reference = reference
|
|
writer._objects[reference.idnum - 1] = stream
|
|
from pypdf.filters import FlateDecode
|
|
from pypdf import get_configuration
|
|
original_decode = FlateDecode.decode
|
|
decodes = []
|
|
def bounded_decode(*args, **kwargs):
|
|
decodes.append(get_configuration().zlib_maximum_output_length)
|
|
return original_decode(*args, **kwargs)
|
|
with patch('engine.ocr_image', side_effect=AssertionError('Must not OCR oversized input')):
|
|
with patch.object(FlateDecode, 'decode', side_effect=bounded_decode):
|
|
self.assertEqual(inspect_content(pdf(mutate=bomb), policy())['decision'], 'block')
|
|
self.assertEqual(decodes, [20 * 1024 * 1024])
|
|
self.assertEqual(inspect_content(pdf(mutate=lambda _w, page: setattr(page.mediabox, 'right', 1000000)), policy())['decision'], 'block')
|
|
|
|
def test_missing_ocr_never_claims_complete(self):
|
|
with self.assertRaises(ValueError): extract_pdf(pdf(), lambda _data: ('', False))
|
|
|
|
def test_debug_stages_record_actual_structure_render_and_ocr(self):
|
|
from stage_timings import collect, metadata, validated
|
|
p = policy(); p['diagnostics'] = {'enabled': True, 'mode': 'debug'}
|
|
with collect(p) as state:
|
|
result = inspect_content(pdf(), p)
|
|
self.assertEqual(result['decision'], 'allow')
|
|
envelope = validated(metadata(state))
|
|
self.assertIsNotNone(envelope)
|
|
stages = {record['stage'] for record in envelope['records']}
|
|
self.assertTrue({'pdf.structure', 'pdf.images', 'pdf.render', 'pdf.ocr'} <= stages)
|
|
for record in envelope['records']:
|
|
self.assertGreaterEqual(record['durationMs'], 0)
|
|
if record['stage'].startswith('pdf.'): self.assertIsNotNone(record['parentId'])
|
|
|
|
|
|
if __name__ == '__main__': unittest.main()
|