pumpwood_communication.encrypt

 13class PumpwoodCryptography:
 14    """Pumpwood data encryption class."""
 15
 16    _fernet_object: Fernet = None
 17    """Fernet object."""
 18
 19    def __init__(self, fernet_key: str = None):
 20        """__init__.
 21
 22        Raises:
 23            PumpWoodOtherException:
 24                Raise PumpWoodOtherException error if Fernet object could not
 25                be correctly configured.
 26        """
 27        if fernet_key is None:
 28            fernet_key = os.getenv('PUMPWOOD_COMUNICATION__CRYPTO_FERNET_KEY')
 29            if fernet_key is None:
 30                log_msg = (
 31                    'PUMPWOOD_COMUNICATION__CRYPTO_FERNET_KEY is not set, '
 32                    'PumpwoodCryptography encription is disable and will '
 33                    'raise error if used')
 34                logger.info(log_msg)
 35            else:
 36                try:
 37                    self._fernet_object = Fernet(fernet_key)
 38                except Exception as e:
 39                    error_str = str(e)
 40                    msg = (
 41                        'Fernet is not correctly cofigured, error when '
 42                        'creating PumpwoodCryptography object. '
 43                        'Error: {erro_msg}')
 44                    raise PumpWoodOtherException(
 45                        message=msg, payload={'erro_msg': error_str})
 46
 47    def _check_if_configured(self) -> bool:
 48        """Check if PumpwoodCryptography is configured.
 49
 50        Returns:
 51            Returns true.
 52
 53        Raises:
 54            PumpWoodOtherException:
 55                Raise PumpWoodOtherException if object was not configured
 56                (Fernet key was not set).
 57        """
 58        if self._fernet_object is None:
 59            msg = (
 60                'PumpwoodCryptography is not configured, fernet key was not '
 61                'passed as argument on object creation and env variable '
 62                '`PUMPWOOD_COMUNICATION__CRYPTO_FERNET_KEY` was not set')
 63            raise PumpWoodOtherException(message=msg)
 64        return True
 65
 66    def encrypt(self, value: any) -> str:
 67        """Encrypt value using object Fernet Key.
 68
 69        It will serialize the data to JSON and if bytes it will be converted to
 70        base64. Data will then be encrypted using the fernet key.
 71
 72        Args:
 73            value (any):
 74                Value that will be encrypted. It is possible to encrypt
 75                python objects and bytes.
 76
 77        Returns:
 78            Return a string with encripted information using objects Fernet
 79            Key.
 80        """
 81        self._check_if_configured()
 82
 83        # None objects should not be encripted
 84        if value is None:
 85            return None
 86
 87        encoded_value = None
 88        # If type is bytes, convert to a dictinary that will be used when
 89        # decrypting data to convert information back to bytes again
 90        if type(value) is bytes:
 91            # Convert bytes to base64 before encription
 92            encoded_base64 = base64.b64encode(value).decode('utf-8')
 93            encoded_value = pumpJsonDump({
 94                '__type__': 'bytes',
 95                '__base64__': encoded_base64,
 96                '__PumpwoodCryptography__': True})
 97        else:
 98            encoded_value = pumpJsonDump(value)
 99        encrypted_data = self._fernet_object.encrypt(encoded_value)
100        return encrypted_data.decode('utf-8')
101
102    def decrypt(self, value: str) -> any:
103        """Decrypt value using object Fernet Key.
104
105        Args:
106            value (str):
107                String value that will be decripted.
108
109        Returns:
110            Return the encrypted object.
111
112        Raise:
113            PumpWoodNotImplementedError:
114                Raise this error if the encrypted data is a dictinary with
115                key `__PumpwoodCryptography__` is equal to `True`, but the
116                key `__type__` indicates
117        """
118        self._check_if_configured()
119        if value is None:
120            return value
121
122        decrypted_value = self._fernet_object.decrypt(value.encode('utf-8'))
123        python_obj = orjson.loads(decrypted_value)
124        if type(python_obj) is not dict:
125            return python_obj
126
127        # Dictionary data with `__PumpwoodCryptography__` key indicates
128        # that the information was treated on encryption
129        is_PumpwoodCryptography = python_obj.get(
130            '__PumpwoodCryptography__', False)
131        if not is_PumpwoodCryptography:
132            return python_obj
133
134        # __type__ will indicated the type of the data that was encrypted,
135        # this will be used to correctly undo the encryption
136        obj_type = python_obj.get('__type__')
137        if obj_type == 'bytes':
138            return base64.b64decode(python_obj['__base64__'])
139
140        msg = (
141            "Object is encrypted using PumpwoodCryptography, but type "
142            "[{obj_type}] could not be recognized. Check if the version of "
143            "Pumpwood Communication used at encryption is compatible")
144        raise PumpWoodNotImplementedError(message=msg, payload={
145            'obj_type': obj_type})