41 lines
2.9 KiB
Python
41 lines
2.9 KiB
Python
|
import io
|
||
|
|
import os
|
||
|
|
import unittest
|
||
|
|
import zipfile
|
||
|
|
from engine import cosine, inspect_content, content_kind
|
||
|
|
|
||
|
|
def policy(): return {"dlp":{"maxBytes":1048576,"rules":[{"id":"secret","kind":"literal","value":"Secret Roadmap"}],"semantic":{"enabled":False}}}
|
||
|
|
|
||
|
|
class InspectionTests(unittest.TestCase):
|
||
|
|
def test_text_and_binary(self):
|
||
|
|
self.assertEqual(inspect_content(b"Public information",policy())["decision"],"allow")
|
||
|
|
self.assertEqual(inspect_content(b"Secret Roadmap",policy())["ruleId"],"secret")
|
||
|
|
self.assertEqual(inspect_content(b"\xff\x00",policy())["decision"],"block")
|
||
|
|
def test_office_extraction(self):
|
||
|
|
output=io.BytesIO()
|
||
|
|
with zipfile.ZipFile(output,"w") as archive:
|
||
|
|
archive.writestr("word/document.xml",'<document xmlns="urn:test"><t>Secret Roadmap</t></document>')
|
||
|
|
result=inspect_content(output.getvalue(),policy())
|
||
|
|
self.assertEqual(result["contentType"],"office");self.assertEqual(result["ruleId"],"secret")
|
||
|
|
def test_pdf_extraction(self):
|
||
|
|
from pypdf import PdfWriter
|
||
|
|
from pypdf.generic import DecodedStreamObject,NameObject,DictionaryObject
|
||
|
|
writer=PdfWriter();page=writer.add_blank_page(300,300)
|
||
|
|
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)})})
|
||
|
|
content=DecodedStreamObject();content.set_data(b'BT /F1 12 Tf 10 100 Td (Secret Roadmap) Tj ET');page[NameObject('/Contents')]=writer._add_object(content)
|
||
|
|
output=io.BytesIO();writer.write(output)
|
||
|
|
result=inspect_content(output.getvalue(),policy());self.assertEqual(result['contentType'],'pdf');self.assertEqual(result['ruleId'],'secret')
|
||
|
|
def test_semantic_threshold_and_missing_model(self):
|
||
|
|
p=policy();p['dlp']['semantic']={'enabled':True,'model':'google/embeddinggemma-2','revision':'test-revision','categories':[{'id':'semantic-secret','threshold':.8,'prototype':[1.,0.]}]}
|
||
|
|
os.environ['EMBEDDING_MODEL_REVISION']='test-revision'
|
||
|
|
result=inspect_content(b'Confidential paraphrase',p,encoder=lambda *_:[1.,0.]);self.assertEqual(result['ruleId'],'semantic-secret')
|
||
|
|
result=inspect_content(b'Confidential paraphrase',p,encoder=lambda *_: (_ for _ in ()).throw(ValueError()));self.assertEqual(result['decision'],'block')
|
||
|
|
def test_cosine_validates(self):
|
||
|
|
self.assertAlmostEqual(cosine([1,0],[1,0]),1)
|
||
|
|
with self.assertRaises(ValueError):cosine([0,0],[0,0])
|
||
|
|
def test_video_has_no_false_safe_coverage(self):
|
||
|
|
video=b'\x00\x00\x00\x18ftypisom'+b'\x00'*20
|
||
|
|
self.assertEqual(content_kind(video),'video');result=inspect_content(video,policy());self.assertEqual(result['decision'],'block');self.assertFalse(result['coverage']['complete'])
|
||
|
|
|
||
|
|
if __name__=='__main__':unittest.main()
|