File size: 2,786 Bytes
a4b70d9 |
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 |
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() |