File size: 3,489 Bytes
931a672 e385db5 931a672 ed6fd56 bcd6d73 931a672 ed6fd56 931a672 ed6fd56 931a672 e385db5 ed6fd56 e385db5 ed6fd56 931a672 aaa82f0 ed6fd56 e385db5 ed6fd56 931a672 ed6fd56 931a672 ed6fd56 e385db5 ed6fd56 e385db5 ed6fd56 e385db5 aaa82f0 e385db5 931a672 ed6fd56 931a672 ed6fd56 931a672 e385db5 ed6fd56 931a672 ed6fd56 e385db5 ed6fd56 931a672 aaa82f0 ed6fd56 931a672 ed6fd56 931a672 ed6fd56 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 |
"""
API Key Generator Script
This script is used to generate API keys for new users.
API keys are generated in the format sk-{username}-{base64_encoded_data}.
The encoded data includes username, creation timestamp, and HMAC signature.
Usage:
python generate_api_key.py <username> <secret_key>
Example:
python generate_api_key.py new_user mock-secret-key-123456789
"""
import base64
import json
import sys
import argparse
import hmac
import hashlib
from datetime import datetime
from loguru import logger
logger.add("logs/api_key_generator.log")
def generate_api_key(username: str, secret_key: str) -> tuple[str, int]:
"""
Generates an API key for the given username using a secret key.
Args:
username (str): Username to generate API key for
secret_key (str): Secret key for API key generation
Returns:
str: Generated API_KEY (ACCESS_KEY, ACCESS_TOKEN)
"""
# Create encoded API key with timestamp
timestamp = int(datetime.now().timestamp())
data = {"username": username, "created_at": timestamp}
json_data = {"username": username, "created_at": timestamp}
json_str = json.dumps(json_data)
logger.debug(f"JSON data for signature: {json_str}")
logger.debug(f"Secret key: {secret_key}")
# Add HMAC signature for additional security
signature = hmac.new(secret_key.encode(), json_str.encode(), hashlib.sha256).hexdigest()
logger.debug(f"Generated signature: {signature}")
data["signature"] = signature
# Create API key
api_key = f"sk-{username}-" + base64.b64encode(json.dumps(data).encode()).decode()
return api_key, timestamp
def save_api_key(username: str, api_key: str, timestamp: int):
"""
Save API key to api_keys.txt file in the same format as the shell script.
Args:
username (str): Username
api_key (str): Generated API key
timestamp (int): Creation timestamp
"""
formatted_timestamp = datetime.fromtimestamp(timestamp).strftime("%Y-%m-%d %H:%M:%S")
with open("api_keys.txt", "a") as f:
f.write(f"Username: {username}\n")
f.write(f"API Key: {api_key}\n")
f.write(f"Timestamp: {formatted_timestamp}\n")
f.write("--------------------------------\n")
def main():
parser = argparse.ArgumentParser(description="API Key Generator")
parser.add_argument("username", help="Username to generate API key for")
parser.add_argument("secret_key", help="Secret key for API key generation")
args = parser.parse_args()
try:
api_key, timestamp = generate_api_key(args.username, args.secret_key)
print("\nAPI Key generated:")
print(f"Username: {args.username}")
print(f"API Key: {api_key}")
# Save to api_keys.txt
save_api_key(args.username, api_key, timestamp)
print("\nAPI Key saved to api_keys.txt")
print("\nUsage example:")
print('curl -X POST "http://localhost:8000/v1/chat/completions" \\')
print(f' -H "Authorization: Bearer {api_key}" \\')
print(' -H "Content-Type: application/json" \\')
print(' -d \'{"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "Hello!"}]}\'')
except argparse.ArgumentError:
print("Please provide a username and secret key", file=sys.stderr)
sys.exit(1)
except Exception as e:
print(f"Error: {str(e)}", file=sys.stderr)
sys.exit(1)
if __name__ == "__main__":
main()
|