89 lines
2.9 KiB
Python
89 lines
2.9 KiB
Python
# 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:
|
|
"""
|
|
Object class for a ssh key
|
|
"""
|
|
def __init__(self, name, key_type, public, private):
|
|
self.name = name
|
|
self.type = key_type
|
|
self.public = public
|
|
self.private = private
|
|
|
|
def get_name(self):
|
|
"""Getter for name"""
|
|
return self.name
|
|
def get_type(self):
|
|
"""Getter for type"""
|
|
return self.type
|
|
def get_private(self):
|
|
"""Getter for private key"""
|
|
return self.private
|
|
def get_public(self):
|
|
"""Getter for public key"""
|
|
return self.public
|
|
|
|
class Collection:
|
|
"""
|
|
Object class of Collection type
|
|
"""
|
|
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):
|
|
"""
|
|
public class to generate a ssh key
|
|
"""
|
|
key = RSA.generate(2048)
|
|
my_ssh_key = SshKey(name=name, key_type=key_type, private=key.exportKey('PEM'), public=key.publickey().exportKey('OpenSSH'))
|
|
self.save_ssh_key(my_ssh_key=my_ssh_key)
|
|
|
|
def save_ssh_key(self, my_ssh_key: SshKey):
|
|
"""
|
|
Function to save the ssh key
|
|
"""
|
|
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()}")
|
|
|
|
## Private Key
|
|
with open(key_file_path, "wb") as private_file:
|
|
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)
|