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