| """ |
| 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) |
| """ |
| |
| 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}") |
|
|
| |
| signature = hmac.new(secret_key.encode(), json_str.encode(), hashlib.sha256).hexdigest() |
|
|
| logger.debug(f"Generated signature: {signature}") |
|
|
| data["signature"] = signature |
|
|
| |
| 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_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() |
|
|
