Add decryption when sshkey are selected

src/encryptor.py

@@ -1,10 +1,12 @@
 import base64
-import os
 from Crypto.Cipher import AES
 from Crypto.Protocol.KDF import PBKDF2
 from Crypto.Random import get_random_bytes
 
 class Encryptor:
+    """
+    Class to encrypt/decrypt content
+    """
     def __init__(self, password: str):
         self.password = password.encode()
         self.salt_size = 16
@@ -13,31 +15,32 @@ class Encryptor:
         self.iterations = 100_000
 
     def _derive_key(self, salt: bytes) -> bytes:
-        """
-        Dérive une clé à partir du mot de passe et du sel.
-        """
         return PBKDF2(self.password, salt, dkLen=self.key_size, count=self.iterations)
 
-    def encrypt(self, plaintext: str) -> str:
+    def encrypt(self, plaintext: str | bytes) -> str:
         """
         Encrypte une chaîne de texte en base64.
         """
+        if isinstance(plaintext, str):
+            plaintext_bytes = plaintext.encode()
+        else:
+            plaintext_bytes = plaintext
+
         salt = get_random_bytes(self.salt_size)
-        key = self._derive_key(salt)
         iv = get_random_bytes(self.iv_size)
+        key = self._derive_key(salt)
 
         # Padding (PKCS7)
-        pad_len = AES.block_size - (len(plaintext.encode()) % AES.block_size)
-        padded = plaintext + chr(pad_len) * pad_len
+        pad_len = AES.block_size - (len(plaintext_bytes) % AES.block_size)
+        padded = plaintext_bytes + bytes([pad_len] * pad_len)
 
         cipher = AES.new(key, AES.MODE_CBC, iv)
-        ciphertext = cipher.encrypt(padded.encode())
+        ciphertext = cipher.encrypt(padded)
 
-        # Encodage final : salt + iv + ciphertext
-        encrypted_data = base64.b64encode(salt + iv + ciphertext).decode()
+        encrypted_data = base64.b64encode(salt + iv + ciphertext)
         return encrypted_data
 
-    def decrypt(self, encrypted_text: str) -> str:
+    def decrypt(self, encrypted_text: str) -> bytes:
         """
         Décrypte une chaîne encodée en base64.
         """
@@ -50,8 +53,7 @@ class Encryptor:
         cipher = AES.new(key, AES.MODE_CBC, iv)
         padded_plaintext = cipher.decrypt(ciphertext)
 
-        # Retrait du padding
+        # Remove padding
         pad_len = padded_plaintext[-1]
-        plaintext = padded_plaintext[:-pad_len].decode()
-
-        return plaintext
+        plaintext = padded_plaintext[:-pad_len]
+        return plaintext  # retourne des bytes