safetynet/apps/inspector/test_pdf_documents.py

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()