Обновлен код, добавлены классы

RSA_Crypto/decrypt.py

@@ -0,0 +1,83 @@
+"""
+Decrypts messages encrypted with RSA.
+"""
+
+import base64
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import serialization, hashes
+from cryptography.hazmat.primitives.asymmetric import padding
+
+
+class RSADecryptor:
+    """
+    Module for decrypting messages encrypted with RSA.
+
+    Provides
+    --------
+
+    *   :class:`RSADecryptor` - class for decrypting messages.
+    """
+    def __init__(self, private_key_file):
+        """
+        Initialize an RSADecryptor with a file containing a private RSA key.
+
+        Parameters
+        ----------
+        private_key_file : str
+            The path to the private key file, which should be in PEM format.
+
+        Attributes
+        ----------
+        private_key_file : str
+            The path to the private key file, which should be in PEM format.
+        private_key : cryptography.hazmat.primitives.asymmetric.rsa.RSAPrivateKey
+            The loaded private key.
+        """
+        self.private_key_file = private_key_file
+        self.private_key = self.load_private_key()
+
+    def load_private_key(self):
+        """
+        Load a private key from a file.
+
+        The file should be in PEM format and contain the private key.
+
+        Returns
+        -------
+        private_key : cryptography.hazmat.primitives.asymmetric.rsa.RSAPrivateKey
+            The loaded private key.
+        """
+        with open(self.private_key_file, 'rb') as f:
+            private_key_data = f.read()
+
+        private_key = serialization.load_pem_private_key(
+            private_key_data,
+            password=None,
+            backend=default_backend()
+        )
+        return private_key
+
+    def decrypt_message(self, encrypted_message):
+        """
+        Decrypt an encrypted message.
+
+        Parameters
+        ----------
+        encrypted_message : str
+            The encrypted message as a base64 encoded string.
+
+        Returns
+        -------
+        decrypted_message : str
+            The decrypted message as a string.
+        """
+        encrypted_bytes = base64.b64decode(encrypted_message)
+        decrypted = self.private_key.decrypt(
+            encrypted_bytes,
+            padding.OAEP(
+                mgf=padding.MGF1(algorithm=hashes.SHA256()),
+                algorithm=hashes.SHA256(),
+                label=None
+            )
+        )
+        return decrypted.decode()

RSA_Crypto/encrypt.py

@@ -0,0 +1,69 @@
+"""
+Encrypt messages encrypted with RSA.
+"""
+
+import base64
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import serialization, hashes
+from cryptography.hazmat.primitives.asymmetric import padding
+
+
+class RSAEncryptor:
+    """
+    Module for encrypting messages encrypted with RSA.
+
+    Provides
+    --------
+
+    *   :class:`RSAEncryptor` - class for encrypting messages.
+    """
+    def __init__(self, public_key_file):
+        self.public_key_file = public_key_file
+        self.public_key = self.load_public_key()
+
+    def load_public_key(self):
+        """
+        Loads a public key from a file.
+
+        Parameters
+        ----------
+        self : `RSAEncryptor`
+            The object itself.
+
+        Returns
+        -------
+        `cryptography.hazmat.primitives.asymmetric.rsa.RSAPublicKey`
+            The loaded key.
+        """
+        with open(self.public_key_file, 'rb') as f:
+            public_key_data = f.read()
+
+        public_key = serialization.load_pem_public_key(
+            public_key_data,
+            backend=default_backend()
+        )
+        return public_key
+
+    def encrypt_message(self, message):
+        """
+        Encrypts a message using RSA public key encryption.
+
+        Parameters
+        ----------
+        message : str
+            The message to be encrypted.
+
+        Returns
+        -------
+        str
+            The encrypted message encoded in Base64.
+        """
+        encrypted = self.public_key.encrypt(
+            message.encode(),
+            padding.OAEP(
+                mgf=padding.MGF1(algorithm=hashes.SHA256()),
+                algorithm=hashes.SHA256(),
+                label=None
+            )
+        )
+        return base64.b64encode(encrypted).decode('utf-8')

RSA_Crypto/generate_save.py

@@ -0,0 +1,104 @@
+"""
+Generate and save RSA key pair
+"""
+
+import os
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
+
+
+class RSAKeyPair:
+    """
+    Class for generating and saving RSA key pair
+    """
+
+    def __init__(self, private_key_file, public_key_file):
+        self.private_key_file = private_key_file
+        self.public_key_file = public_key_file
+        self.private_key = None
+        self.public_key = None
+
+    def generate_private_key(self):
+        """
+        Generate RSA private key.
+
+        The generated key has a 2048 bit length and a public exponent of 65537.
+        """
+        self.private_key = rsa.generate_private_key(
+            public_exponent=65537,
+            key_size=2048,
+            backend=default_backend()
+        )
+
+    def generate_public_key(self):
+        """
+        Generate the public key from the private key.
+
+        Raises:
+            ValueError: Private key must be generated first.
+        """
+        if self.private_key is None:
+            raise ValueError("Private key must be generated first.")
+        self.public_key = self.private_key.public_key()
+
+    def save_private_key(self):
+        """
+        Save the generated private key to a file.
+
+        The private key is saved in PEM format with the Traditional OpenSSL
+        format and no encryption. If the private key has not been generated yet,
+        a ValueError is raised.
+
+        Raises:
+            ValueError: Private key must be generated before saving.
+        """
+        if self.private_key is None:
+            raise ValueError("Private key must be generated before saving.")
+        with open(self.private_key_file, 'wb') as f:
+            f.write(self.private_key.private_bytes(
+                encoding=serialization.Encoding.PEM,
+                format=serialization.PrivateFormat.TraditionalOpenSSL,
+                encryption_algorithm=serialization.NoEncryption()
+            ))
+
+    def save_public_key(self):
+        """
+        Save the generated public key to a file.
+
+        The public key is saved in PEM format with the SubjectPublicKeyInfo
+        format. If the public key has not been generated yet, a ValueError is
+        raised.
+
+        Raises:
+            ValueError: Public key must be generated before saving.
+        """
+        if self.public_key is None:
+            raise ValueError("Public key must be generated before saving.")
+        with open(self.public_key_file, 'wb') as f:
+            f.write(self.public_key.public_bytes(
+                encoding=serialization.Encoding.PEM,
+                format=serialization.PublicFormat.SubjectPublicKeyInfo
+            ))
+
+    def generate_and_save_keys(self):
+        """
+        Generate and save the private and public keys to the specified files.
+
+        This method first checks if the key files already exist. If not, it
+        generates the private key and saves it to the private key file, and
+        generates the public key based on the private key and saves it to the
+        public key file. If the files already exist, it simply prints a message
+        to the console.
+
+        This method does not raise any exceptions if the files already exist.
+        """
+        if not os.path.exists(self.private_key_file) and not os.path.exists(self.public_key_file):
+            try:
+                self.generate_private_key()
+                self.save_private_key()
+                self.generate_public_key()
+                self.save_public_key()
+            except ValueError as e:
+                print(e)
+

main.py

@@ -1,7 +1,16 @@
+"""
+    This is a simple RSA encryptor/decryptor.
+    It uses a key pair which is stored in a 'keys' directory.
+    The keys are generated when you run the script.
+    The script will encrypt a message for you if you provide it.
+    The script will decrypt a message for you if you provide it.
+    The script will generate a key pair and save it in the 'keys' directory.
+"""
+
 import os
-from utilits.encrypt import encrypt_message
+from RSA_Crypto.encrypt import RSAEncryptor
-from utilits.decrypt import decrypt_message
+from RSA_Crypto.decrypt import RSADecryptor
-from utilits.generate_save import generate_private_key, generate_public_key, save_private_key, save_public_key
+from RSA_Crypto.generate_save import RSAKeyPair
 
 KEYS_DIR = 'keys'
 PRIVATE_KEY_FILE = f'{KEYS_DIR}/private-key.pem'
@@ -10,24 +19,16 @@ PUBLIC_KEY_FILE = f'{KEYS_DIR}/public-key.pem'
 if not os.path.exists(KEYS_DIR):
     os.makedirs(KEYS_DIR)
 
-if not os.path.exists(PRIVATE_KEY_FILE) and not os.path.exists(PUBLIC_KEY_FILE):
+key_pair = RSAKeyPair(PRIVATE_KEY_FILE, PUBLIC_KEY_FILE,)
-    private_key = generate_private_key()
+key_pair.generate_and_save_keys()
-    save_private_key(private_key, PRIVATE_KEY_FILE)
-    public_key = generate_public_key(private_key)
-    save_public_key(public_key, PUBLIC_KEY_FILE)
-    print(f"Новые ключи сгенерированы и сохранены в '{PRIVATE_KEY_FILE}' и '{PUBLIC_KEY_FILE}'.")
-else:
-    if os.path.exists(PRIVATE_KEY_FILE) and os.path.exists(PUBLIC_KEY_FILE):
-        pass
-    else:
-        print(f"Файлы ключей не найдены. Пожалуйста, сгенерируйте новые ключи.")
-        exit()
 
 # Шифрование сообщения
 message = "Это лицензионный ключ."
-encrypted_message = encrypt_message(message, PUBLIC_KEY_FILE)
+rsa_encryptor = RSAEncryptor(PUBLIC_KEY_FILE)
+encrypted_message = rsa_encryptor.encrypt_message(message)
 print("Зашифрованное сообщение (Base64):", encrypted_message)
 
 # Дешифрование сообщения
-decrypted_message = decrypt_message(encrypted_message, PRIVATE_KEY_FILE)
+rsa_decryptor = RSADecryptor(PRIVATE_KEY_FILE)
+decrypted_message = rsa_decryptor.decrypt_message(encrypted_message)
 print("Расшифрованное сообщение:", decrypted_message)

utilits/decrypt.py

@@ -1,30 +0,0 @@
-import base64
-from cryptography.hazmat.backends import default_backend
-from cryptography.hazmat.primitives import serialization, hashes
-from cryptography.hazmat.primitives.asymmetric import padding
-
-
-def load_keys_private(private_key_file):
-    with open(private_key_file, 'rb') as f:
-        private_key_data = f.read()
-
-    private_key = serialization.load_pem_private_key(
-        private_key_data,
-        password=None,
-        backend=default_backend()
-    )
-    return private_key
-
-
-def decrypt_message(encrypted_message, private_key_file):
-    private_key = load_keys_private(private_key_file)
-    encrypted_bytes = base64.b64decode(encrypted_message)
-    decrypted = private_key.decrypt(
-        encrypted_bytes,
-        padding.OAEP(
-            mgf=padding.MGF1(algorithm=hashes.SHA256()),
-            algorithm=hashes.SHA256(),
-            label=None
-        )
-    )
-    return decrypted.decode()

utilits/encrypt.py

@@ -1,28 +0,0 @@
-import base64
-from cryptography.hazmat.backends import default_backend
-from cryptography.hazmat.primitives import serialization, hashes
-from cryptography.hazmat.primitives.asymmetric import padding
-
-
-def load_keys_public(public_key_file):
-    with open(public_key_file, 'rb') as f:
-        public_key_data = f.read()
-
-    public_key = serialization.load_pem_public_key(
-        public_key_data,
-        backend=default_backend()
-    )
-    return public_key
-
-
-def encrypt_message(message, public_key_file):
-    public_key = load_keys_public(public_key_file)
-    encrypted = public_key.encrypt(
-        message.encode(),
-        padding.OAEP(
-            mgf=padding.MGF1(algorithm=hashes.SHA256()),
-            algorithm=hashes.SHA256(),
-            label=None
-        )
-    )
-    return base64.b64encode(encrypted).decode('utf-8')

utilits/generate_save.py

@@ -1,34 +0,0 @@
-from cryptography.hazmat.backends import default_backend
-from cryptography.hazmat.primitives import serialization
-from cryptography.hazmat.primitives.asymmetric import rsa
-
-
-def generate_private_key():
-    private_key = rsa.generate_private_key(
-        public_exponent=65537,
-        key_size=2048,
-        backend=default_backend()
-    )
-    return private_key
-
-
-def generate_public_key(private_key):
-    public_key = private_key.public_key()
-    return public_key
-
-
-def save_private_key(private_key, private_key_file):
-    with open(private_key_file, 'wb') as f:
-        f.write(private_key.private_bytes(
-            encoding=serialization.Encoding.PEM,
-            format=serialization.PrivateFormat.TraditionalOpenSSL,
-            encryption_algorithm=serialization.NoEncryption()
-        ))
-
-
-def save_public_key(public_key, public_key_file):
-    with open(public_key_file, 'wb') as f:
-        f.write(public_key.public_bytes(
-            encoding=serialization.Encoding.PEM,
-            format=serialization.PublicFormat.SubjectPublicKeyInfo
-        ))