Inspect static PDF and OOXML content with bounded local coverage

apps/inspector/test_pdf_documents.py

@@ -0,0 +1,117 @@
+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()