Initial server source import
This commit is contained in:
sashatrask
@@ -0,0 +1,168 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user