xltdp/python3/efhandle.py

import base64
import os
import json
import yaml
from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.primitives import serialization, hashes
from cryptography.exceptions import InvalidSignature

class EFHAMDLE:
    def __init__(self, private_key: rsa.RSAPrivateKey = None, public_key: rsa.RSAPublicKey = None):
        """Initializes the EFHAMDLE class.

        Arguments:
        private_key -- the private key to use for decryption and signing (optional)
        public_key -- the public key to use for encryption and signature verification (optional)
        """
        self.private_key = private_key
        self.public_key = public_key

    def encrypt_string_to_file(self, string: str, file: str):
        """Encrypts a string and writes it to a file.

        Arguments:
        string -- the string to be encrypted
        file -- the file to write the encrypted string to
        """
        # Encrypt the string using the public key
        encrypted_string = self.public_key.encrypt(
            string.encode(),
            padding.OAEP(
                mgf=padding.MGF1(algorithm=hashes.SHA256()),
                algorithm=hashes.SHA256(),
                label=None
            )
        )

        # Write the encrypted string to the file
        with open(file, "wb") as f:
            f.write(encrypted_string)

    def decrypt_string_from_file(self, file: str) -> str:
        """Decrypts a string from a file.

        Arguments:
        file -- the file to read the encrypted string from

        Returns:
        The decrypted string.
        """
        # Read the encrypted string from the file
        with open(file, "rb") as f:
            encrypted_string = f.read()

        # Decrypt the string using the private key
        return self.private_key.decrypt(
            encrypted_string,
            padding.OAEP(
                mgf=padding.MGF1(algorithm=hashes.SHA256()),
                algorithm=hashes.SHA256(),
                label=None
            )
        ).decode()

    @staticmethod
    def generate_rsa_key_pair():
        """Generates a new RSA key pair.

        Returns:
        A tuple containing the private key and public key.
        """
        private_key = rsa.generate_private_key(
            public_exponent=65537,
            key_size=2048
        )
        public_key = private_key.public_key()
        return private_key, public_key

    @staticmethod
    def save_rsa_key_pair(private_key: rsa.RSAPrivateKey, public_key: rsa.RSAPublicKey, private_key_file: str, public_key_file: str):

   """Saves an RSA key pair to files.

        Arguments:
        private_key -- the private key to be saved
        public_key -- the public key to be saved
        private_key_file -- the file to save the private key to
        public_key_file -- the file to save the public key to
        """
        with open(private_key_file, "wb") as f:
            f.write(private_key.private_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PrivateFormat.PKCS8,
                encryption_algorithm=serialization.NoEncryption()
            ))
        with open(public_key_file, "wb") as f:
            f.write(public_key.public_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PublicFormat.SubjectPublicKeyInfo
            ))

    @staticmethod
    def load_rsa_private_key(private_key_file: str) -> rsa.RSAPrivateKey:
        """Loads an RSA private key from a file.

        Arguments:
        private_key_file -- the file to load the private key from

        Returns:
        The private key.
        """
        with open(private_key_file, "rb") as f:
            private_key = serialization.load_pem_private_key(
                f.read(),
                password=None,
                backend=default_backend()
            )
        return private_key

    @staticmethod
    def load_rsa_public_key(public_key_file: str) -> rsa.RSAPublicKey:
        """Loads an RSA public key from a file.

        Arguments:
        public_key_file -- the file to load the public key from

        Returns:
        The public key.
        """
        with open(public_key_file, "rb") as f:
            public_key = serialization.load_pem_public_key(
                f.read(),
                backend=default_backend()
            )
        return public_key

    def sign_data(self, data: bytes) -> bytes:
        """Signs data using an RSA private key.

        Arguments:
        data -- the data to be signed

        Returns:
        The signature.
        """
        return self.private_key.sign(
            data,
            padding.PSS(
                mgf=padding.MGF1(hashes.SHA256()),
                salt_length=padding.PSS.MAX_LENGTH
            ),
            hashes.SHA256()
        )

    def verify_signature(self, data: bytes, signature: bytes) -> bool:
        """Verifies the signature of data using an RSA public key.

        Arguments:
        data -- the data that was signed
        signature -- the signature to be verified

        Returns:
        True if the signature is valid, False otherwise.
        """
        try:
            self.public_key.verify(
                signature,
                data,
                padding.PSS(
                    mgf=padding.MGF1(hashes.SHA256()),
                    salt_length=padding.PSS.MAX_LENGTH
                ),
                hashes.SHA256()
            )
            return True
        except InvalidSignature:
            return False

    def encrypt_json_to_file(self, json_data: dict, file: str):
        """Encrypts a JSON object and writes it to a file.

        Arguments:
        json_data -- the JSON object to be encrypted
        file -- the file to write the encrypted JSON object to
        """
        # Convert the JSON object to a string
        json_string = json.dumps(json_data)

        # Encrypt the string and write it to the file
        self.encrypt_string_to_file(json_string, file + ".ejson")

    def encrypt_yaml_to_file(self, yaml_data: dict, file: str):
        """Encrypts a YAML object and writes it to a file.

        Arguments:
        yaml_data -- the YAML object to be encrypted
        file -- the file to write the encrypted YAML object to
        """
        # Convert the YAML object to a string
        yaml_string = yaml.dump(yaml_data)

        # Encrypt the string and write it to the file
        self.encrypt_string_to_file(yaml_string, file + ".yamle")