Add decryption when sshkey are selected

src/collection.py

@@ -1,6 +1,9 @@
+# pylint: disable=line-too-long, C0114
 from pathlib import Path
 from os import chmod
 from Crypto.PublicKey import RSA
+import yaml
+from encryptor import Encryptor
 
 class SshKey:
     """
@@ -29,9 +32,10 @@ class Collection:
     """
     Object class of Collection type
     """
-    def __init__(self, collection_name):
+    def __init__(self, collection_name: str, password: str):
         self.collection_name = collection_name
         self.collection_path = Path.home().joinpath(".sshkeymanager", self.collection_name)
+        self.encryptor = Encryptor(password)
 
     def generate_ssh_key(self, name: str, key_type: str):
         """
@@ -48,13 +52,37 @@ class Collection:
         key_file_path = self.collection_path.joinpath(my_ssh_key.get_name())
         ## Info File
         with open(f"{key_file_path}.txt", "w+", encoding="utf-8") as info_file:
-            info_file.write(f"name: {my_ssh_key.get_name()}\nKey_type: {my_ssh_key.get_type()}")
+            info_file.write(f"name: {my_ssh_key.get_name()}\nkey_type: {my_ssh_key.get_type()}")
 
         ## Private Key
         with open(key_file_path, "wb") as private_file:
-            private_file.write(my_ssh_key.get_private())
+            encrypted_key = self.encryptor.encrypt(my_ssh_key.get_private())
+            private_file.write(encrypted_key)
             chmod(key_file_path, 0o600)
 
         ## Public Key
         with open(f"{key_file_path}.pub", "wb") as public_file:
             public_file.write(my_ssh_key.get_public())
+
+    def get_ssh_key(self, name: str) -> SshKey:
+        """
+        Get ssh key and decrypt private key
+        """
+        key_file_path = self.collection_path.joinpath(name)
+
+        # Info file
+        with open(f"{key_file_path}.txt", "r", encoding="utf-8") as info_file:
+            data = yaml.safe_load(info_file)
+            name = data["name"]
+            key_type = data["key_type"]
+
+        # Private Key
+        with open(key_file_path, "rb") as private_file:
+            encrypted_private_key = private_file.read()
+            private_key = self.encryptor.decrypt(encrypted_private_key)
+
+        # Public key
+        with open(f"{key_file_path}.pub", "rb") as public_file:
+            public_key = public_file.read()
+
+        return SshKey(name=name, key_type=key_type, private=private_key, public=public_key)