Source code for schrodinger.application.matsci.enc

"""
Utilities for encryption.

Copyright Schrodinger, LLC. All rights reserved.
"""

import os

from cryptography.fernet import Fernet

from schrodinger.utils import fileutils

ENC_KEYS_PATH = os.path.join(fileutils.get_mmshare_data_dir(), '.enc_keys')
ENC_FILE_EXT = '.enc'

CG_KEY_FILE_NAME = 'cg'


def get_encryption_key(key_file_name, key_file_path=None):
    """
    Return the encryption key for the given file name.  If the
    file doesn't exist then an encryption key will be created for
    it.

    :type key_file_name: str
    :param key_file_name: the key file name

    :type key_file_path: str or None
    :param key_file_path: the key file path, if None then the
        default Schrodinger path is used

    :rtype: bytes
    :return: the encryption key
    """

    if key_file_path is None:
        key_file_path = ENC_KEYS_PATH
    key_file_path = os.path.join(key_file_path, key_file_name)

    if os.path.exists(key_file_path):
        with open(key_file_path, 'rb') as f_handle:
            key = f_handle.read()
    else:
        key = Fernet.generate_key()
        with open(key_file_path, 'wb') as f_handle:
            f_handle.write(key)

    return key


def write_encrypted_file(in_file_path, key_file_name, key_file_path=None):
    """
    Write an encrypted version of the given input file path.

    :type in_file_path: str
    :param in_file_path: the in file path to encrypt

    :type key_file_name: str
    :param key_file_name: the key file name to use in the encryption

    :type key_file_path: str or None
    :param key_file_path: the key file path, if None then the
        default Schrodinger path is used

    :raise ValueError: if there is a problem

    :rtype: str
    :return: the encrypted file path
    """

    if in_file_path.endswith(ENC_FILE_EXT):
        raise ValueError('encrypting an encrypted file is not supported')

    with open(in_file_path) as f_handle:
        text = f_handle.read()

    return write_encrypted_text(text,
                                in_file_path,
                                key_file_name,
                                key_file_path=key_file_path)


def write_encrypted_text(text, in_file_path, key_file_name, key_file_path=None):
    """
    Write the unencrypted text to an encrypted file with the given
    path.

    :type text: str
    :param text: the text to encrypt

    :type in_file_path: str
    :param in_file_path: the file path to which to write the encrypted
        text

    :type key_file_name: str
    :param key_file_name: the key file name to use in the encryption

    :type key_file_path: str or None
    :param key_file_path: the key file path, if None then the
        default Schrodinger path is used

    :rtype: str
    :return: the encrypted file path
    """

    key = get_encryption_key(key_file_name, key_file_path=key_file_path)
    f_obj = Fernet(key)
    token = f_obj.encrypt(text.encode())

    out_file_path = in_file_path
    if not out_file_path.endswith(ENC_FILE_EXT):
        out_file_path += ENC_FILE_EXT

    with open(out_file_path, 'wb') as f_handle:
        f_handle.write(token)

    return out_file_path


def get_unencrypted_text(in_file_path, key_file_name, key_file_path=None):
    """
    Return the unencrypted text for the given encrypted input file path.

    :type in_file_path: str
    :param in_file_path: the encrypted in file path to unencrypt

    :type key_file_name: str
    :param key_file_name: the key file name to use in the unencryption

    :type key_file_path: str or None
    :param key_file_path: the key file path, if None then the
        default Schrodinger path is used

    :raise ValueError: if there is a problem

    :rtype: str
    :return: the unencrypted text
    """

    if not in_file_path.endswith(ENC_FILE_EXT):
        raise ValueError('the input file is not encrypted')

    with open(in_file_path, 'rb') as f_handle:
        token = f_handle.read()

    key = get_encryption_key(key_file_name, key_file_path=key_file_path)
    f_obj = Fernet(key)
    text = f_obj.decrypt(token).decode()

    return text


def write_unencrypted_file(in_file_path, key_file_name, key_file_path=None):
    """
    Write an unencrypted version for the given encrypted input file path.

    :type in_file_path: str
    :param in_file_path: the encrypted in file path to unencrypt

    :type key_file_name: str
    :param key_file_name: the key file name to use in the unencryption

    :type key_file_path: str or None
    :param key_file_path: the key file path, if None then the
        default Schrodinger path is used

    :rtype: str
    :return: the unencrypted file path
    """

    text = get_unencrypted_text(in_file_path,
                                key_file_name,
                                key_file_path=key_file_path)
    out_file_path = fileutils.splitext(in_file_path)[0]
    with open(out_file_path, 'w') as f_handle:
        f_handle.write(text)

    return out_file_path