Source code for schrodinger.application.matsci.enc

"""
Utilities for encryption.

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

import os
import pymmlibs

from cryptography.fernet import Fernet

from schrodinger import structure
from schrodinger.utils import fileutils
from schrodinger.infra import mm
from schrodinger.application.desmond import cms
from schrodinger.application.desmond import ffiostructure

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

CG_KEY_FILE_NAME = 'cg'

FST_ENC_STR_KEY = f's{mm.M2IO_BACKEND_HIDDEN_PROP_PATTERN}fst_enc_str'

FFIO_FUNCT_KEY = 's_ffio_funct'
INVALID = 'INVALID'


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


def _get_enc_str_from_fst(fst):
    """
    Get an encrypted string representation of the given FFIOStructure.

    :type fst: ffiostructure.FFIOStructure
    :param fst: the structure

    :rtype: str
    :return: the encrypted string representation
    """

    fst_str = ffiostructure.write_cms_to_string(fst)
    return pymmlibs._encode_string(fst_str)


def _get_fst_from_enc_str(enc_str):
    """
    Get an FFIOStructure from the given encrypted string representation.

    :type enc_str: str
    :param enc_str: the encrypted string representation

    :rtype: ffiostructure.FFIOStructure
    :return: the structure
    """

    fst_str = pymmlibs._decode_string(enc_str)
    # encode/decode doesn't work properly, at a minimum
    # the following newline is needed
    fst_str += '\n'
    with structure.StructureReader.fromString(fst_str) as reader:
        st = next(reader)
    return ffiostructure.FFIOStructure(st)


def _get_prop_fst_from_fst(fst):
    """
    Get an FFIOStructure with a propertized encrypted string representation.

    :type fst: ffiostructure.FFIOStructure
    :param fst: the structure

    :rtype: ffiostructure.FFIOStructure
    :return: the structure with the property
    """

    fst = fst.copy()
    enc_str = _get_enc_str_from_fst(fst)
    fst.property[FST_ENC_STR_KEY] = enc_str
    return fst


def _get_fst_from_prop_fst(fst):
    """
    Get an FFIOStructure from a propertized encrypted string representation.

    :type fst: ffiostructure.FFIOStructure
    :param fst: the structure with the property

    :rtype: ffiostructure.FFIOStructure
    :return: the structure
    """

    fst = fst.copy()
    enc_str = fst.property.pop(FST_ENC_STR_KEY)
    return _get_fst_from_enc_str(enc_str)


def _hide_sub_block_properties(sub_block_property):
    """
    Hide sub-block properties in the given _FFIOSubBlockProperty.

    :type sub_block_property: ffiostructure._FFIOSubBlockProperty
    :param sub_block_property: the sub-block properties
    """

    for key in sub_block_property.keys():
        if key.startswith('r'):
            sub_block_property[key] = 0
        elif key == FFIO_FUNCT_KEY:
            # doesn't have to be consistent but use
            # the OPLS convention anyway
            sub_block_property[key] = INVALID


def _hide_ffio_ff(fst):
    """
    Get a copy of the given FFIOStructure where data
    in all ffio_ff sub-blocks has been hidden.

    :type fst: ffiostructure.FFIOStructure
    :param fst: the structure

    :rtype: ffiostructure.FFIOStructure
    :return: the structure with the data hidden
    """

    fst = fst.copy()
    # for some reason the following must be done like this
    for obj in fst.ffio.site:
        _hide_sub_block_properties(obj.property)
    for obj in fst.ffio.bond:
        _hide_sub_block_properties(obj.property)
    for obj in fst.ffio.angle:
        _hide_sub_block_properties(obj.property)
    for obj in fst.ffio.dihedral:
        _hide_sub_block_properties(obj.property)
    for obj in fst.ffio.exclusion:
        _hide_sub_block_properties(obj.property)
    for obj in fst.ffio.pair:
        _hide_sub_block_properties(obj.property)
    for obj in fst.ffio.vdwtype:
        _hide_sub_block_properties(obj.property)
    for obj in fst.ffio.vdwtypescombined:
        _hide_sub_block_properties(obj.property)
    for obj in fst.ffio.restraint:
        _hide_sub_block_properties(obj.property)
    for obj in fst.ffio.virtual:
        _hide_sub_block_properties(obj.property)
    for obj in fst.ffio.pseudo:
        _hide_sub_block_properties(obj.property)
    for obj in fst.ffio.constraint:
        _hide_sub_block_properties(obj.property)
    for obj in fst.ffio.posfbhw:
        _hide_sub_block_properties(obj.property)
    for obj in fst.ffio.anglefbhw:
        _hide_sub_block_properties(obj.property)
    for obj in fst.ffio.improperfbhw:
        _hide_sub_block_properties(obj.property)
    for obj in fst.ffio.stretchfbhw:
        _hide_sub_block_properties(obj.property)
    return fst


def _encrypt_fst(fst):
    """
    Get an encrypted FFIOStructure from the given one.

    :type fst: ffiostructure.FFIOStructure
    :param fst: the structure

    :rtype: ffiostructure.FFIOStructure
    :return: the encrypted structure
    """

    prop_fst = _get_prop_fst_from_fst(fst)
    enc_fst = _hide_ffio_ff(prop_fst)
    return enc_fst


def _decrypt_fst(enc_fst):
    """
    Get an FFIOStructure from the given encrypted one.

    :type enc_fst: ffiostructure.FFIOStructure
    :param enc_fst: the encrypted structure

    :rtype: ffiostructure.FFIOStructure
    :return: the structure
    """

    return _get_fst_from_prop_fst(enc_fst)


def encrypt_cms(in_cms_fp, out_cms_fp=None):
    """
    Encrypt the FF parameters in the given cms file.

    :type in_cms_fp: str
    :param in_cms_fp: the file path for the cms to be encrypted

    :type out_cms_fp: str or None
    :param out_cms_fp: the file path to which to write the encrypted cms
        if None then the input file path will be overwritten
    """

    out_cms_fp = out_cms_fp or in_cms_fp

    cms_obj = cms.Cms(in_cms_fp)

    st = cms_obj._raw_fsys_ct
    st.write(out_cms_fp)

    for comp_ct in cms_obj.comp_ct:
        fst = ffiostructure.FFIOStructure(comp_ct)
        enc_fst = _encrypt_fst(fst)
        ffiostructure.write_cms(enc_fst, out_cms_fp, mode=mm.M2IO_APPEND)


def decrypt_cms(in_cms_fp, out_cms_fp=None):
    """
    Decrypt the FF parameters in the given cms file.

    :type in_cms_fp: str
    :param in_cms_fp: the file path for the cms to be decrypted

    :type out_cms_fp: str or None
    :param out_cms_fp: the file path to which to write the decrypted cms
        if None then the input file path will be overwritten
    """

    out_cms_fp = out_cms_fp or in_cms_fp

    cms_obj = cms.Cms(in_cms_fp)

    st = cms_obj._raw_fsys_ct
    st.write(out_cms_fp)

    for comp_ct in cms_obj.comp_ct:
        enc_fst = ffiostructure.FFIOStructure(comp_ct)
        fst = _decrypt_fst(enc_fst)
        ffiostructure.write_cms(fst, out_cms_fp, mode=mm.M2IO_APPEND)