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