168 lines
5.3 KiB
Python
168 lines
5.3 KiB
Python
import json
|
|
import os
|
|
import secrets
|
|
import time
|
|
|
|
from utils.redis_utils import get_redis_client
|
|
|
|
# Redis에 API 키를 저장할 때 사용할 접두사
|
|
API_KEY_PREFIX = "api_key:"
|
|
# Docker 컨테이너의 /workspace 디렉토리에 파일을 저장하도록 절대 경로 사용
|
|
API_KEYS_FILE = "/workspace/api_keys.json"
|
|
|
|
|
|
def _read_keys_from_file():
|
|
"""Helper function to read all keys from the JSON file."""
|
|
if not os.path.exists(API_KEYS_FILE):
|
|
return {}
|
|
with open(API_KEYS_FILE, "r") as f:
|
|
try:
|
|
return json.load(f)
|
|
except json.JSONDecodeError:
|
|
return {}
|
|
|
|
|
|
def _write_keys_to_file(keys):
|
|
"""Helper function to write all keys to the JSON file."""
|
|
with open(API_KEYS_FILE, "w") as f:
|
|
json.dump(keys, f, indent=4)
|
|
|
|
|
|
import redis
|
|
|
|
|
|
def load_api_keys_from_file():
|
|
"""
|
|
JSON 파일에서 API 키를 읽어 Redis에 로드합니다.
|
|
Redis 연결 실패 시 몇 초간 재시도하여 시작 시점의 문제를 해결합니다.
|
|
"""
|
|
keys_from_file = _read_keys_from_file()
|
|
if not keys_from_file:
|
|
print("API key file not found or empty. Skipping loading.")
|
|
return
|
|
|
|
redis_client = get_redis_client()
|
|
max_retries = 5
|
|
retry_delay = 2 # 초
|
|
|
|
for i in range(max_retries):
|
|
try:
|
|
# Redis 연결 테스트
|
|
redis_client.ping()
|
|
|
|
# 연결 성공 시 키 로드
|
|
for key_name, key_data in keys_from_file.items():
|
|
if not redis_client.exists(key_name):
|
|
redis_client.hset(key_name, mapping=key_data)
|
|
print(f"Loaded API key from file: {key_name}")
|
|
|
|
print("Successfully loaded all keys into Redis.")
|
|
return # 성공 시 함수 종료
|
|
|
|
except redis.exceptions.ConnectionError as e:
|
|
print(f"Could not connect to Redis (attempt {i+1}/{max_retries}): {e}")
|
|
if i < max_retries - 1:
|
|
print(f"Retrying in {retry_delay} seconds...")
|
|
time.sleep(retry_delay)
|
|
else:
|
|
print("Failed to load API keys into Redis after multiple retries.")
|
|
break
|
|
|
|
|
|
def generate_api_key(prefix="sk") -> str:
|
|
"""안전한 API 키를 생성합니다. (예: sk-xxxxxxxx)"""
|
|
return f"{prefix}-{secrets.token_hex(16)}"
|
|
|
|
|
|
def create_api_key(client_name: str, key_prefix="sk") -> dict:
|
|
"""
|
|
새로운 API 키를 생성하고 Redis와 파일에 저장합니다.
|
|
"""
|
|
api_key = generate_api_key(prefix=key_prefix)
|
|
redis_client = get_redis_client()
|
|
|
|
key_storage_name = f"{API_KEY_PREFIX}{api_key}"
|
|
key_data = {
|
|
"client_name": client_name,
|
|
"created_at": str(int(time.time())),
|
|
"is_active": "true",
|
|
}
|
|
|
|
# Redis에 저장 (hset 사용)
|
|
redis_client.hset(key_storage_name, mapping=key_data)
|
|
|
|
# 파일에 즉시 저장
|
|
all_keys = _read_keys_from_file()
|
|
all_keys[key_storage_name] = key_data
|
|
_write_keys_to_file(all_keys)
|
|
|
|
return {"api_key": api_key, **key_data}
|
|
|
|
|
|
def validate_api_key(api_key: str) -> bool:
|
|
"""
|
|
제공된 API 키가 유효한지 검증합니다. decode_responses=True로 인해 모든 값은 문자열입니다.
|
|
1. Redis에서 먼저 확인합니다.
|
|
2. Redis에 없으면 api_keys.json 파일에서 확인합니다.
|
|
3. 파일에서 유효한 키를 찾으면 Redis에 다시 동기화합니다.
|
|
"""
|
|
if not api_key:
|
|
return False
|
|
|
|
redis_client = get_redis_client()
|
|
key_storage_name = f"{API_KEY_PREFIX}{api_key}"
|
|
|
|
# 1. Redis에서 확인 (decode_responses=True이므로 반환값은 문자열)
|
|
is_active_in_redis = redis_client.hget(key_storage_name, "is_active")
|
|
if is_active_in_redis == "true":
|
|
return True
|
|
|
|
# 2. Redis에 없으면 파일에서 확인
|
|
all_keys_from_file = _read_keys_from_file()
|
|
key_data_from_file = all_keys_from_file.get(key_storage_name)
|
|
|
|
if key_data_from_file and key_data_from_file.get("is_active") == "true":
|
|
# 3. 파일에 유효한 키가 있으면 Redis에 다시 기록 (Self-healing, hset 사용)
|
|
redis_client.hset(key_storage_name, mapping=key_data_from_file)
|
|
print(f"Key '{key_storage_name}' not found in Redis, but restored from file.")
|
|
return True
|
|
|
|
return False
|
|
|
|
|
|
def revoke_api_key(api_key: str) -> bool:
|
|
"""
|
|
API 키를 Redis와 파일에서 삭제하여 폐기합니다.
|
|
"""
|
|
redis_client = get_redis_client()
|
|
key_storage_name = f"{API_KEY_PREFIX}{api_key}"
|
|
|
|
# Redis에서 삭제
|
|
result = redis_client.delete(key_storage_name)
|
|
|
|
if result > 0:
|
|
# 파일에서도 삭제
|
|
all_keys = _read_keys_from_file()
|
|
if key_storage_name in all_keys:
|
|
del all_keys[key_storage_name]
|
|
_write_keys_to_file(all_keys)
|
|
return True
|
|
return False
|
|
|
|
|
|
def list_api_keys() -> list:
|
|
"""
|
|
저장된 모든 API 키의 목록을 반환합니다.
|
|
(주의: 실제 환경에서는 키 자체를 노출하지 않는 것이 좋습니다)
|
|
"""
|
|
redis_client = get_redis_client()
|
|
keys = []
|
|
|
|
# decode_responses=True이므로 모든 키와 값은 문자열.
|
|
for key_name in redis_client.scan_iter(f"{API_KEY_PREFIX}*"):
|
|
key_data = redis_client.hgetall(key_name)
|
|
key_data["api_key"] = key_name.replace(API_KEY_PREFIX, "", 1)
|
|
keys.append(key_data)
|
|
|
|
return keys
|