diff --git a/shared_utils/key_pattern_manager.py b/shared_utils/key_pattern_manager.py index 8feb0640..138a3b09 100644 --- a/shared_utils/key_pattern_manager.py +++ b/shared_utils/key_pattern_manager.py @@ -8,13 +8,20 @@ from shared_utils.config_loader import get_conf as get_conf pj = os.path.join default_user_name = 'default_user' - +# match openai keys +openai_regex = re.compile( + r"sk-[a-zA-Z0-9_-]{48}$|" + + r"sk-[a-zA-Z0-9_-]{92}$|" + + r"sk-proj-[a-zA-Z0-9_-]{48}$|"+ + r"sk-proj-[a-zA-Z0-9_-]{124}$|"+ + r"sess-[a-zA-Z0-9]{40}$" +) def is_openai_api_key(key): CUSTOM_API_KEY_PATTERN = get_conf('CUSTOM_API_KEY_PATTERN') if len(CUSTOM_API_KEY_PATTERN) != 0: API_MATCH_ORIGINAL = re.match(CUSTOM_API_KEY_PATTERN, key) else: - API_MATCH_ORIGINAL = re.match(r"sk-[a-zA-Z0-9]{48}$|sk-proj-[a-zA-Z0-9]{48}$|sess-[a-zA-Z0-9]{40}$", key) + API_MATCH_ORIGINAL = openai_regex.match(key) return bool(API_MATCH_ORIGINAL) diff --git a/tests/test_key_pattern_manager.py b/tests/test_key_pattern_manager.py index bf84441b..b555d8b3 100644 --- a/tests/test_key_pattern_manager.py +++ b/tests/test_key_pattern_manager.py @@ -21,10 +21,32 @@ class TestKeyPatternManager(unittest.TestCase): key = "sx-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx" self.assertFalse(is_openai_api_key(key)) - key = "sess-wg61ZafYHpNz7FFwIH7HGZlbVqUVaeV5tatHCWpl" + key = "sess-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx" self.assertTrue(is_openai_api_key(key)) - key = "sess-wg61ZafYHpNz7FFwIH7HGZlbVqUVa5tatHCWpl" + key = "sess-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx" + self.assertFalse(is_openai_api_key(key)) + + key = "sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx_xxxxxxxxxxxxxxxxxx-xxxxxxxxxxxxxxxxxxxx" + self.assertTrue(is_openai_api_key(key)) + key = "sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx-xxxxxxxxxxxxxxx_xxxxxxxxxxxx_xxxxx-xxxxxxxxxxxxxxxxxxxx" + self.assertTrue(is_openai_api_key(key)) + key = "sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx-xxx-xxxxxxxxxxxxxxx_xxxxxxxxxxxx_xxxxx-xxxxxx-xxxxxxxxxxxxx" + self.assertTrue(is_openai_api_key(key)) + key = "sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx-xxx-xxxxxxxxxxxxxxx_xxxxxxxxxxxx_xxxxx-xxxxxxxxxxxxxxxxxx" + self.assertFalse(is_openai_api_key(key)) + key = "sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx-xxx-xxxxxxxxxxxxxxx_xxxxxxxxxxxx_xxxxx-xxxxxxxxxxxxxxxxxxxxx" + self.assertFalse(is_openai_api_key(key)) + + key = "sk-proj-xx-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx_xxxxxxxxxxxxxxxxxxxxx-xxxxxxxx" + self.assertTrue(is_openai_api_key(key)) + key = "sk-proj-xx-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx_xxxxxxxxxxxx_xxxxxxxxxxxxxxxxxxxxx-xxxxxxxx" + self.assertTrue(is_openai_api_key(key)) + key = "sk-proj-xx-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx_xxxxxxxxxxxx_xxxxxxxxxxxxxxxxxx-xxxxxxxx" + self.assertFalse(is_openai_api_key(key)) + key = "sk-proj-xx-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx_xxxxxxxxxxxx_xxxxxxxxxxxxxxxxxx-xxxxxxxxxxxxx" + self.assertFalse(is_openai_api_key(key)) + key = "sk-proj-xx-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx_xxxxxxxxxxxx_xxxxxxxxxxxxxxxxxx-xxx啊xxxxxxx" self.assertFalse(is_openai_api_key(key))