File size: 4,741 Bytes
b8989d2 1e6eac0 |
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 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 |
# agents/tools/crypto.py
from cryptography.hazmat.primitives.asymmetric import rsa, ed25519
from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.kdf.scrypt import Scrypt
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.backends import default_backend
import os
import base64
backend = default_backend()
# ========== 🔐 RSA KEYS ==========
def generate_rsa_key_pair(key_size=2048):
private_key = rsa.generate_private_key(
public_exponent=65537,
key_size=key_size,
backend=backend
)
public_key = private_key.public_key()
return private_key, public_key
# ========== 🌀 ED25519 KEYS ==========
def generate_ed25519_key_pair():
private_key = ed25519.Ed25519PrivateKey.generate()
public_key = private_key.public_key()
return private_key, public_key
# ========== 💾 SERIALIZATION ==========
def serialize_private_key(private_key, password: str = None):
if password:
encryption_algorithm = serialization.BestAvailableEncryption(password.encode())
else:
encryption_algorithm = serialization.NoEncryption()
return private_key.private_bytes(
encoding=serialization.Encoding.PEM,
format=serialization.PrivateFormat.PKCS8,
encryption_algorithm=encryption_algorithm
)
def serialize_public_key(public_key):
return public_key.public_bytes(
encoding=serialization.Encoding.PEM,
format=serialization.PublicFormat.SubjectPublicKeyInfo
)
def load_private_key(pem_data, password: str = None):
return serialization.load_pem_private_key(
pem_data,
password=password.encode() if password else None,
backend=backend
)
def load_public_key(pem_data):
return serialization.load_pem_public_key(
pem_data,
backend=backend
)
def generate_keypair(method="rsa", password: bytes = None):
"""
Создаёт пару ключей (приватный, публичный) с заданным методом.
method: "rsa" или "ed25519"
password: если указан, приватный ключ будет зашифрован
Возвращает (private_key_pem: bytes, public_key_pem: bytes)
"""
if method == "rsa":
private_key = rsa.generate_private_key(
public_exponent=65537, key_size=2048, backend=default_backend()
)
elif method == "ed25519":
private_key = ed25519.Ed25519PrivateKey.generate()
else:
raise ValueError("Unsupported key generation method")
encryption_algorithm = (
serialization.BestAvailableEncryption(password)
if password
else serialization.NoEncryption()
)
private_pem = private_key.private_bytes(
encoding=serialization.Encoding.PEM,
format=serialization.PrivateFormat.PKCS8,
encryption_algorithm=encryption_algorithm,
)
public_pem = private_key.public_key().public_bytes(
encoding=serialization.Encoding.PEM,
format=serialization.PublicFormat.SubjectPublicKeyInfo,
)
return private_pem, public_pem
# ========== 🔐 ENCRYPT / DECRYPT PRIVATE KEY BY SYMMETRIC KEY ==========
def derive_key(password: str, salt: bytes = None):
if not salt:
salt = os.urandom(16)
kdf = Scrypt(
salt=salt,
length=32,
n=2**14,
r=8,
p=1,
backend=backend
)
key = kdf.derive(password.encode())
return key, salt
def encrypt_data(data: bytes, password: str):
key, salt = derive_key(password)
aesgcm = AESGCM(key)
nonce = os.urandom(12)
encrypted = aesgcm.encrypt(nonce, data, None)
return {
'ciphertext': base64.b64encode(encrypted).decode(),
'salt': base64.b64encode(salt).decode(),
'nonce': base64.b64encode(nonce).decode()
}
def decrypt_data(encrypted_data: dict, password: str):
ciphertext = base64.b64decode(encrypted_data['ciphertext'])
salt = base64.b64decode(encrypted_data['salt'])
nonce = base64.b64decode(encrypted_data['nonce'])
key, _ = derive_key(password, salt=salt)
aesgcm = AESGCM(key)
return aesgcm.decrypt(nonce, ciphertext, None)
# ========== ✅ TESTING ==========
if __name__ == "__main__":
priv, pub = generate_ed25519_key_pair()
priv_pem = serialize_private_key(priv)
pub_pem = serialize_public_key(pub)
print("PRIVATE PEM:")
print(priv_pem.decode())
print("PUBLIC PEM:")
print(pub_pem.decode())
encrypted = encrypt_data(priv_pem, "secret-password")
decrypted = decrypt_data(encrypted, "secret-password")
assert decrypted == priv_pem
print("✅ Encryption/decryption OK")
|