diff --git a/g4f/api/__init__.py b/g4f/api/__init__.py index 75536691..adb1243b 100644 --- a/g4f/api/__init__.py +++ b/g4f/api/__init__.py @@ -82,7 +82,7 @@ from .stubs import ( from g4f import debug try: - from g4f.gui.server.crypto import create_or_read_keys, decrypt_data + from g4f.gui.server.crypto import create_or_read_keys, decrypt_data, get_session_key has_crypto = True except ImportError: has_crypto = False @@ -217,6 +217,7 @@ class Api: print(f"Register authentication key: {''.join(['*' for _ in range(len(AppConfig.g4f_api_key))])}") if has_crypto: private_key, _ = create_or_read_keys() + session_key = get_session_key() @self.app.middleware("http") async def authorization(request: Request, call_next): if AppConfig.g4f_api_key is not None or AppConfig.demo: @@ -231,10 +232,15 @@ class Api: if has_crypto and user_g4f_api_key: try: expires, user = decrypt_data(private_key, user_g4f_api_key).split(":", 1) - expires = int(expires) - int(time.time()) - debug.log(f"User: '{user}' G4F API key expires in {expires} seconds") - except Exception as e: - return ErrorResponse.from_message(f"Invalid G4F API key: {e}", HTTP_401_UNAUTHORIZED) + except: + try: + data = json.loads(decrypt_data(session_key, user_g4f_api_key)) + expires = int(decrypt_data(private_key, data["data"])) + 86400 + user = data.get("user", None) + except: + return ErrorResponse.from_message(f"Invalid G4F API key", HTTP_401_UNAUTHORIZED) + expires = int(expires) - int(time.time()) + debug.log(f"User: '{user}' G4F API key expires in {expires} seconds") if expires < 0: return ErrorResponse.from_message("G4F API key expired", HTTP_401_UNAUTHORIZED) else: diff --git a/g4f/gui/server/backend_api.py b/g4f/gui/server/backend_api.py index db11f037..91fcd982 100644 --- a/g4f/gui/server/backend_api.py +++ b/g4f/gui/server/backend_api.py @@ -29,7 +29,7 @@ try: except ImportError as e: has_markitdown = False try: - from .crypto import rsa, serialization, create_or_read_keys, decrypt_data, encrypt_data + from .crypto import rsa, serialization, create_or_read_keys, decrypt_data, encrypt_data, get_session_key has_crypto = True except ImportError: has_crypto = False @@ -79,7 +79,7 @@ class Backend_Api(Api): self.chat_cache = {} if has_crypto: - private_key_obj = rsa.generate_private_key(public_exponent=65537, key_size=4096) + private_key_obj = get_session_key() public_key_obj = private_key_obj.public_key() public_key_pem = public_key_obj.public_bytes( encoding=serialization.Encoding.PEM, diff --git a/g4f/gui/server/crypto.py b/g4f/gui/server/crypto.py index 976bd0eb..a909bde8 100644 --- a/g4f/gui/server/crypto.py +++ b/g4f/gui/server/crypto.py @@ -9,6 +9,15 @@ from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicKey, RSAPriva from ...cookies import get_cookies_dir +SESSION_KEY: RSAPrivateKey = None + +def get_session_key() -> RSAPrivateKey: + global SESSION_KEY + if SESSION_KEY is not None: + return SESSION_KEY + SESSION_KEY = rsa.generate_private_key(public_exponent=65537, key_size=4096) + return SESSION_KEY + def create_or_read_keys() -> tuple[RSAPrivateKey, RSAPublicKey]: private_key_file = os.path.join(get_cookies_dir(), "private_key.pem") public_key_file = os.path.join(get_cookies_dir(), "public_key.pem")