mirror of
https://github.com/xtekky/gpt4free.git
synced 2025-10-29 18:51:59 +08:00
82 lines
2.7 KiB
Python
82 lines
2.7 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import base64
|
|
|
|
from cryptography.hazmat.primitives import serialization
|
|
from cryptography.hazmat.primitives.asymmetric import padding, rsa
|
|
from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicKey, RSAPrivateKey
|
|
|
|
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")
|
|
|
|
if os.path.isfile(private_key_file) and os.path.isfile(public_key_file):
|
|
# Read private key
|
|
with open(private_key_file, 'rb') as f:
|
|
private_key_pem = f.read()
|
|
private_key = serialization.load_pem_private_key(
|
|
private_key_pem,
|
|
password=None # Use password=b'mypassword' here if the key is encrypted
|
|
)
|
|
|
|
# Read public key
|
|
with open(public_key_file, 'rb') as f:
|
|
public_key_pem = f.read()
|
|
public_key = serialization.load_pem_public_key(public_key_pem)
|
|
|
|
return private_key, public_key
|
|
|
|
# Generate keys
|
|
private_key_obj = rsa.generate_private_key(public_exponent=65537, key_size=1024)
|
|
public_key_obj = private_key_obj.public_key()
|
|
|
|
# Serialize private key
|
|
private_key_pem = private_key_obj.private_bytes(
|
|
encoding=serialization.Encoding.PEM,
|
|
format=serialization.PrivateFormat.PKCS8,
|
|
encryption_algorithm=serialization.NoEncryption()
|
|
)
|
|
|
|
# Serialize public key
|
|
public_key_pem = public_key_obj.public_bytes(
|
|
encoding=serialization.Encoding.PEM,
|
|
format=serialization.PublicFormat.SubjectPublicKeyInfo
|
|
)
|
|
|
|
os.makedirs(os.path.dirname(private_key_file), exist_ok=True)
|
|
|
|
# Write the private key PEM to a file
|
|
with open(private_key_file, 'wb') as f:
|
|
f.write(private_key_pem)
|
|
|
|
# Write the public key PEM to a file
|
|
with open(public_key_file, 'wb') as f:
|
|
f.write(public_key_pem)
|
|
|
|
return private_key_obj, public_key_obj
|
|
|
|
def decrypt_data(private_key_obj: RSAPrivateKey, encrypted_data: str) -> str:
|
|
decrypted = private_key_obj.decrypt(
|
|
base64.b64decode(encrypted_data),
|
|
padding.PKCS1v15()
|
|
)
|
|
return decrypted.decode()
|
|
|
|
def encrypt_data(public_key: RSAPublicKey, decrypted_data: str) -> str:
|
|
encrypted = public_key.encrypt(
|
|
decrypted_data.encode(),
|
|
padding.PKCS1v15()
|
|
)
|
|
return base64.b64encode(encrypted).decode() |