68 lines
4.5 KiB
Python
68 lines
4.5 KiB
Python
|
import unittest
|
||
|
|
from unittest.mock import patch
|
||
|
|
import io
|
||
|
|
from engine import scan_download, reference_document, inspect_content
|
||
|
|
|
||
|
|
def policy(source='rule malicious { strings: $a = "MALWARE_TEST_PAYLOAD" condition: $a }'):
|
||
|
|
return {'security':{'yara':{'enabled':True,'rules':[{'id':'central-malware','source':source}], 'maxScanBytes':1048576}},
|
||
|
|
'dlp':{'maxBytes':1048576,'rules':[],'semantic':{'enabled':False}}}
|
||
|
|
|
||
|
|
class SecurityTests(unittest.TestCase):
|
||
|
|
def test_real_yara_match(self):
|
||
|
|
import yara
|
||
|
|
result=scan_download(b'MZ\x00MALWARE_TEST_PAYLOAD',policy())
|
||
|
|
self.assertEqual(result['decision'],'block')
|
||
|
|
self.assertEqual(result['ruleId'],'central-malware')
|
||
|
|
self.assertTrue(result['coverage']['complete'])
|
||
|
|
def test_binary_malware_coverage_separate_from_dlp(self):
|
||
|
|
data=b'MZ\x00\xffbenign'
|
||
|
|
self.assertEqual(scan_download(data,policy())['decision'],'allow')
|
||
|
|
self.assertEqual(inspect_content(data,policy())['decision'],'block')
|
||
|
|
def test_invalid_or_empty_rules_fail_closed(self):
|
||
|
|
self.assertEqual(scan_download(b'safe',policy('invalid'))['decision'],'block')
|
||
|
|
p=policy();p['security']['yara']['rules']=[]
|
||
|
|
self.assertEqual(scan_download(b'safe',p)['decision'],'block')
|
||
|
|
self.assertEqual(scan_download(b'safe',{})['decision'],'block')
|
||
|
|
def test_include_and_console_import_rejected(self):
|
||
|
|
self.assertEqual(scan_download(b'safe',policy('include "C:/secret"'))['decision'],'block')
|
||
|
|
self.assertEqual(scan_download(b'safe',policy('import "console" rule leak { condition: console.log("secret") }'))['decision'],'block')
|
||
|
|
def test_size_failure(self):
|
||
|
|
p=policy();p['security']['yara']['maxScanBytes']=2
|
||
|
|
self.assertEqual(scan_download(b'large',p)['ruleId'],'yara-size-limit')
|
||
|
|
def test_reference_registration_contains_no_plaintext(self):
|
||
|
|
text=b'Internal confidential roadmap for new product launch during next quarter'
|
||
|
|
result=reference_document(text)
|
||
|
|
self.assertTrue(result['registrationReady']);self.assertTrue(result['fragmentHashes'])
|
||
|
|
self.assertNotIn('confidential',str(result))
|
||
|
|
self.assertFalse(reference_document(b'\xff\x00')['registrationReady'])
|
||
|
|
def test_training_truncation_never_claims_complete(self):
|
||
|
|
result=reference_document(' '.join('distinctword'+str(index) for index in range(1200)).encode())
|
||
|
|
self.assertEqual(len(result['fragmentHashes']),1000)
|
||
|
|
self.assertGreater(result['coverage']['fragmentHashesTotal'],1000)
|
||
|
|
self.assertTrue(result['coverage']['fragmentHashesTruncated'])
|
||
|
|
self.assertTrue(result['coverage']['extractionComplete'])
|
||
|
|
self.assertFalse(result['coverage']['complete'])
|
||
|
|
self.assertFalse(result['registrationReady'])
|
||
|
|
def test_text_training_does_not_create_temporary_content_files(self):
|
||
|
|
with patch('engine.tempfile.TemporaryDirectory',side_effect=AssertionError('Raw content persistence forbidden')):
|
||
|
|
result=reference_document(b'Confidential strategic roadmap contains enough words for fingerprint training')
|
||
|
|
self.assertTrue(result['registrationReady'])
|
||
|
|
def test_trained_fragments_block_reformatted_excerpt(self):
|
||
|
|
text=b'Internal confidential roadmap for new product launch during next quarter'
|
||
|
|
reference=reference_document(text)
|
||
|
|
p=policy();p['dlp']['rules']=[{'id':'trained','kind':'fragment-sha256','value':value} for value in reference['fragmentHashes']]
|
||
|
|
result=inspect_content(b'Introduction. INTERNAL confidential roadmap for new product launch during next quarter. Ending.',p)
|
||
|
|
self.assertEqual(result['ruleId'],'trained')
|
||
|
|
def test_real_bundled_ocr_offline_detects_image_secret(self):
|
||
|
|
from PIL import Image, ImageDraw, ImageFont
|
||
|
|
image=Image.new('RGB',(1000,180),'white')
|
||
|
|
ImageDraw.Draw(image).text((30,45),'Secret Roadmap',font=ImageFont.truetype('C:/Windows/Fonts/arial.ttf',60),fill='black')
|
||
|
|
output=io.BytesIO();image.save(output,format='PNG')
|
||
|
|
p=policy();p['dlp']['rules']=[{'id':'image-secret','kind':'literal','value':'Secret Roadmap'}]
|
||
|
|
# Inference and lazy initialization must not enter a model-download path.
|
||
|
|
with patch('rapidocr.utils.download_file.DownloadFile.run',side_effect=AssertionError('Network model download forbidden')):
|
||
|
|
result=inspect_content(output.getvalue(),p)
|
||
|
|
self.assertEqual(result['ruleId'],'image-secret')
|
||
|
|
self.assertTrue(result['coverage']['complete'])
|
||
|
|
|
||
|
|
if __name__=='__main__':unittest.main()
|