0e20e4375a
增强API密钥强度检测算法，综合评估字符集熵值和实际字符多样性，避免低复杂度密钥通过验证。同时改进解密功能的错误处理，统一异常类型。 主要改进： - 新增字符多样性检测，防止全重复字符的弱密钥通过验证 - 修正熵值计算逻辑，基于原始字符集而非密钥实际字符分布 - 增强解密函数异常处理，统一抛出InvalidToken异常类型 - 更新测试用例以匹配新的强校验逻辑 技术细节： - check_key_strength函数新增min_diversity、charset、exclude_ambiguous参数 - 解密时捕获UnicodeDecodeError并转换为InvalidToken - 测试代码配合新逻辑调整断言和参数传递
121 lines
4.5 KiB
Python
121 lines
4.5 KiB
Python
import pytest
|
||
import string
|
||
from llmapikeygen import (
|
||
generate_api_key, generate_batch, check_key_strength,
|
||
encrypt_key, decrypt_key, PYDANTIC_AVAILABLE, CRYPTOGRAPHY_AVAILABLE
|
||
)
|
||
|
||
# 跳过整个文件（如果依赖未安装）
|
||
pytestmark = [
|
||
pytest.mark.skipif(not PYDANTIC_AVAILABLE, reason="Pydantic not installed"),
|
||
pytest.mark.skipif(not CRYPTOGRAPHY_AVAILABLE, reason="Cryptography not installed"),
|
||
]
|
||
|
||
|
||
def test_generate_api_key_pydantic_validation():
|
||
"""测试 pydantic 参数校验"""
|
||
# 长度小于最小值（8）
|
||
with pytest.raises(ValueError) as excinfo:
|
||
generate_api_key(length=7)
|
||
assert "Length must be ≥ 8" in str(excinfo.value) or "至少8个字符" in str(excinfo.value)
|
||
|
||
# 前缀为空字符串
|
||
with pytest.raises(ValueError) as excinfo:
|
||
generate_api_key(prefix="")
|
||
assert "Prefix cannot be empty" in str(excinfo.value) or "前缀不能为空" in str(excinfo.value)
|
||
|
||
# 前缀仅含空格
|
||
with pytest.raises(ValueError) as excinfo:
|
||
generate_api_key(prefix=" ")
|
||
assert "Prefix cannot be empty" in str(excinfo.value) or "前缀不能为空" in str(excinfo.value)
|
||
|
||
|
||
def test_generate_api_key_use_crypto():
|
||
"""测试使用 cryptography 增强随机性"""
|
||
# 单密钥生成
|
||
key1 = generate_api_key(length=32, use_crypto=True)
|
||
assert len(key1) == len("sk-") + 32
|
||
# 批量生成（确保不重复）
|
||
batch = generate_batch(count=10, use_crypto=True)
|
||
assert len(set(batch)) == 10
|
||
# 不同字符集下的兼容性
|
||
key_charset = generate_api_key(charset=string.ascii_lowercase, use_crypto=True)
|
||
assert all(c in string.ascii_lowercase for c in key_charset[3:])
|
||
|
||
|
||
def test_check_key_strength_crypto_enhanced():
|
||
"""测试 cryptography 增强的密钥强度检测"""
|
||
# 生成高强度密钥
|
||
key = generate_api_key(length=64, use_crypto=True)
|
||
is_secure, msg = check_key_strength(key)
|
||
assert is_secure
|
||
assert "合格" in msg
|
||
|
||
# 边缘情况：熵值刚好达标（修改部分：传递charset和exclude_ambiguous参数）
|
||
charset = string.ascii_letters[:16] # 原始字符集：a-p（16种）
|
||
# 生成时关闭exclude_ambiguous，避免排除o字符（确保字符集大小=16）
|
||
key_marginal = generate_api_key(
|
||
charset=charset,
|
||
exclude_ambiguous=False, # 关闭易混淆字符排除
|
||
length=32
|
||
)
|
||
is_secure, msg = check_key_strength(
|
||
key_marginal,
|
||
min_entropy=4.0,
|
||
charset=charset, # 传递原始字符集
|
||
exclude_ambiguous=False # 与生成时一致
|
||
)
|
||
assert is_secure
|
||
|
||
|
||
def test_encrypt_decrypt_key_basic():
|
||
"""测试密钥加密和解密"""
|
||
# 普通密钥加密解密
|
||
key = generate_api_key()
|
||
encrypted_key, secret_key = encrypt_key(key)
|
||
# 断言加密后的数据非空
|
||
assert isinstance(encrypted_key, bytes) and len(encrypted_key) > 0
|
||
assert isinstance(secret_key, bytes) and len(secret_key) > 0
|
||
# 解密后与原密钥一致
|
||
decrypted_key = decrypt_key(encrypted_key, secret_key)
|
||
assert decrypted_key == key
|
||
|
||
|
||
def test_encrypt_decrypt_custom_secret_key():
|
||
"""测试自定义加密密钥"""
|
||
from cryptography.fernet import Fernet
|
||
# 自定义 secret_key（需符合 Fernet 格式）
|
||
custom_secret_key = Fernet.generate_key()
|
||
key = "sk-custom-1234567890abcdef"
|
||
# 用自定义密钥加密
|
||
encrypted_key, secret_key = encrypt_key(key, secret_key=custom_secret_key)
|
||
assert secret_key == custom_secret_key # 返回的密钥与自定义一致
|
||
# 解密
|
||
decrypted_key = decrypt_key(encrypted_key, custom_secret_key)
|
||
assert decrypted_key == key
|
||
|
||
|
||
def test_encrypt_decrypt_special_key():
|
||
"""测试含特殊字符的密钥加密解密"""
|
||
# 含特殊符号的密钥
|
||
special_key = "sk-!@#$%^&*()_+-=[]{}|;:,.?~`1234"
|
||
encrypted_key, secret_key = encrypt_key(special_key)
|
||
decrypted_key = decrypt_key(encrypted_key, secret_key)
|
||
assert decrypted_key == special_key
|
||
|
||
# 长密钥（128位主体）
|
||
long_key = generate_api_key(length=128)
|
||
encrypted_long, secret_long = encrypt_key(long_key)
|
||
decrypted_long = decrypt_key(encrypted_long, secret_long)
|
||
assert decrypted_long == long_key
|
||
|
||
|
||
def test_generate_batch_use_crypto():
|
||
"""测试批量生成（启用 cryptography）"""
|
||
batch = generate_batch(count=5, length=24, use_crypto=True)
|
||
assert len(batch) == 5
|
||
for key in batch:
|
||
assert len(key) == len("sk-") + 24
|
||
# 强度检测合格
|
||
is_secure, _ = check_key_strength(key)
|
||
assert is_secure |