truf-server/tests/test_custom_provider_detector_compatibility.py

import hashlib
import json
import os
from pathlib import Path
import re
import shutil
import subprocess
import sys
import tempfile
import unittest

import yaml


ROOT = Path(__file__).resolve().parents[1]
APP_DIR = ROOT / 'app'
POLICY_PATH = APP_DIR / 'trufflehog-custom-detectors.yaml'
sys.path.insert(0, str(APP_DIR))

from keycheck_candidates import extract_candidates
import scanner


def synthetic_material(label, length):
    alphabet = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789'
    seed = hashlib.sha512(label.encode('ascii')).hexdigest()
    output = ''
    while len(output) < length:
        output += ''.join(
            alphabet[int(seed[index:index + 2], 16) % len(alphabet)]
            for index in range(0, len(seed), 2)
        )
        seed = hashlib.sha512(seed.encode('ascii')).hexdigest()
    return output[:length]


def configured_trufflehog():
    candidates = [
        os.getenv('TRUF_TEST_TRUFFLEHOG'),
        r'C:\Tools\trufflehog.exe',
        shutil.which('trufflehog'),
    ]
    return next((Path(value) for value in candidates if value and Path(value).is_file()), None)


class CustomProviderDetectorPolicyTests(unittest.TestCase):
    def test_context_alternatives_are_independent_and_canonicalized(self):
        policy = yaml.safe_load(POLICY_PATH.read_text(encoding='utf-8'))
        detectors = {item['name']: item for item in policy['detectors']}
        expected = {
            'Xai': 'xai_context_before',
            'XaiContextAfter': 'xai_context_after',
            'ZaiGLM': 'zai_glm_context_before',
            'ZaiGLMContextAfter': 'zai_glm_context_after',
        }

        for name, regex_name in expected.items():
            with self.subTest(name=name):
                self.assertEqual(list(detectors[name]['regex']), [regex_name])

        findings = [
            {
                'DetectorName': 'CustomRegex',
                'ExtraData': {'name': 'XaiContextAfter'},
            },
            {
                'DetectorName': 'CustomRegex',
                'ExtraData': {'name': 'ZaiGLMContextAfter'},
            },
        ]
        scanner.normalize_custom_detector_names(findings)

        self.assertEqual(
            [finding['DetectorName'] for finding in findings], ['Xai', 'ZaiGLM'],
        )
        self.assertTrue(all(
            finding['OriginalDetectorName'] == 'CustomRegex' for finding in findings
        ))

    def test_both_context_directions_match_and_route_without_the_cli(self):
        policy = yaml.safe_load(POLICY_PATH.read_text(encoding='utf-8'))
        detectors = {item['name']: item for item in policy['detectors']}
        xai_key = 'xai-' + synthetic_material('xai-source-policy', 48)
        zai_key = ('a' * 32) + '.' + synthetic_material('zai-source-policy', 16)
        cases = (
            ('xai-before', f'XAI_API_KEY={xai_key}', xai_key,
             'Xai', 'XaiContextAfter', 'Xai', 'xai'),
            ('xai-after', f'{xai_key} api.x.ai', xai_key,
             'XaiContextAfter', 'Xai', 'Xai', 'xai'),
            ('zai-before', f'ZAI_API_KEY={zai_key}', zai_key,
             'ZaiGLM', 'ZaiGLMContextAfter', 'ZaiGLM', 'zai'),
            ('zai-after', f'{zai_key} api.z.ai', zai_key,
             'ZaiGLMContextAfter', 'ZaiGLM', 'ZaiGLM', 'zai'),
        )

        for label, content, key, emitted, opposite, canonical, service in cases:
            with self.subTest(label=label):
                regex = next(iter(detectors[emitted]['regex'].values()))
                match = re.search(regex, content)
                self.assertIsNotNone(match)
                self.assertEqual(match.group(1), key)
                opposite_regex = next(iter(detectors[opposite]['regex'].values()))
                self.assertIsNone(re.search(opposite_regex, content))

                finding = {
                    'DetectorName': 'CustomRegex', 'Raw': key,
                    'ExtraData': {'name': emitted},
                }
                scanner.normalize_custom_detector_names([finding])
                self.assertEqual(finding['DetectorName'], canonical)
                self.assertEqual(finding['OriginalDetectorName'], 'CustomRegex')
                self.assertEqual(
                    [candidate.service for candidate in extract_candidates(finding)],
                    [service],
                )


@unittest.skipUnless(configured_trufflehog(), 'configured TruffleHog binary is unavailable')
class CustomProviderDetectorCLITests(unittest.TestCase):
    def test_real_cli_detects_both_context_directions_and_routes_candidates(self):
        executable = configured_trufflehog()
        xai_key = 'xai-' + synthetic_material('xai-custom-compatibility', 48)
        zai_key = 'zai-' + synthetic_material('zai-custom-compatibility', 48)
        cases = {
            'xai-before': (f'XAI_API_KEY={xai_key}', 'Xai', 'Xai', 'xai'),
            'xai-after': (f'{xai_key} api.x.ai', 'XaiContextAfter', 'Xai', 'xai'),
            'zai-before': (f'ZAI_API_KEY={zai_key}', 'ZaiGLM', 'ZaiGLM', 'zai'),
            'zai-after': (f'{zai_key} api.z.ai', 'ZaiGLMContextAfter', 'ZaiGLM', 'zai'),
        }

        with tempfile.TemporaryDirectory() as temp_dir:
            for name, (content, emitted_name, canonical_name, service) in cases.items():
                with self.subTest(name=name):
                    fixture = Path(temp_dir) / f'{name}.env'
                    fixture.write_text(content + '\n', encoding='utf-8', newline='\n')
                    completed = subprocess.run(
                        [
                            str(executable), 'filesystem', str(fixture),
                            '--config', str(POLICY_PATH), '--json', '--no-update',
                            '--no-verification',
                            '--results', 'verified,unknown,unverified,filtered_unverified',
                        ],
                        stdin=subprocess.DEVNULL,
                        stdout=subprocess.PIPE,
                        stderr=subprocess.PIPE,
                        check=False,
                        timeout=60,
                    )
                    self.assertEqual(completed.returncode, 0)
                    findings = []
                    for line in completed.stdout.splitlines():
                        value = json.loads(line.decode('utf-8', errors='strict'))
                        if value.get('DetectorName') == 'CustomRegex':
                            findings.append(value)

                    self.assertEqual(len(findings), 1)
                    self.assertEqual(findings[0]['ExtraData']['name'], emitted_name)
                    scanner.normalize_custom_detector_names(findings)
                    self.assertEqual(findings[0]['DetectorName'], canonical_name)
                    self.assertEqual(findings[0]['OriginalDetectorName'], 'CustomRegex')
                    self.assertEqual(
                        [candidate.service for candidate in extract_candidates(findings[0])],
                        [service],
                    )


if __name__ == '__main__':
    unittest.main()