github.com/sentienttechnologies/studio-go-runner@v0.0.0-20201118202441-6d21f2ced8ee/assets/crypto/encryptor.py (about)

     1  # Copyright 2020 (c) Cognizant Digital Business, Evolutionary AI. All rights reserved. Issued under the Apache 2.0 License.
     2  #
     3  # Credit to Andrei Denissov, Cognizant, April 2020
     4  import os
     5  import sys
     6  from Crypto.PublicKey import RSA
     7  from Crypto.Cipher import PKCS1_OAEP
     8  from Crypto.Hash import SHA256
     9  import nacl.secret
    10  import nacl.utils
    11  import base64
    12  
    13  
    14  class Encryptor:
    15      """
    16      Implementation for experiment payload builder
    17      using public key RSA encryption.
    18      """
    19      def __init__(self, keypath: str):
    20          """
    21          param: keypath - file path to .pem file with public key
    22          """
    23  
    24          self.key_path = os.path.abspath(keypath)
    25          self.recipient_key = None
    26          try:
    27              self.recipient_key = RSA.import_key(open(self.key_path).read())
    28          except:
    29              print(
    30                  "FAILED to import recipient public key : {0}".format(self.key_path))
    31              return
    32  
    33      def _import_rsa_key(self, key_path: str):
    34          key = None
    35          try:
    36              key = RSA.import_key(open(key_path).read())
    37          except:
    38              self.logger.error(
    39                  "FAILED to import RSA key from: {0}".format(key_path))
    40              key = None
    41          return key
    42  
    43      def _encrypt_str(self, workload: str):
    44          # Generate one-time symmetric session key:
    45          session_key = nacl.utils.random(32)
    46  
    47          # Encrypt the data with the NaCL session key
    48          data_to_encrypt = workload.encode("utf-8")
    49          box_out = nacl.secret.SecretBox(session_key)
    50          encrypted_data = box_out.encrypt(data_to_encrypt)
    51          encrypted_data_text = base64.b64encode(encrypted_data)
    52  
    53          # Encrypt the session key with the public RSA key
    54          cipher_rsa = PKCS1_OAEP.new(key=self.recipient_key, hashAlgo=SHA256)
    55          encrypted_session_key = cipher_rsa.encrypt(session_key)
    56          encrypted_session_key_text = base64.b64encode(encrypted_session_key)
    57  
    58          return encrypted_session_key_text, encrypted_data_text
    59  
    60      def _decrypt_data(self, private_key_path, encrypted_key_text,
    61                        encrypted_data_text):
    62          private_key = self._import_rsa_key(private_key_path)
    63          if private_key is None:
    64              return None
    65  
    66          try:
    67              private_key = RSA.import_key(open(private_key_path).read())
    68          except:
    69              self.logger.error(
    70                  "FAILED to import private key from: {0}".format(private_key_path))
    71              return None
    72  
    73          # Decrypt the session key with the private RSA key
    74          cipher_rsa = PKCS1_OAEP.new(key=private_key, hashAlgo=SHA256)
    75          session_key = cipher_rsa.decrypt(
    76              base64.b64decode(encrypted_key_text))
    77  
    78          # Decrypt the data with the NaCL session key
    79          box_in = nacl.secret.SecretBox(session_key)
    80          decrypted_data = box_in.decrypt(
    81              base64.b64decode(encrypted_data_text))
    82          decrypted_data = decrypted_data.decode("utf-8")
    83  
    84          return decrypted_data
    85  
    86      def encrypt(self, payload: str):
    87          enc_key, enc_payload = self._encrypt_str(payload)
    88  
    89          enc_key_str = enc_key.decode("utf-8")
    90          enc_payload_str = enc_payload.decode("utf-8")
    91  
    92          return "{0},{1}".format(enc_key_str, enc_payload_str)
    93  
    94  
    95  def main():
    96      if len(sys.argv) < 4:
    97          print("USAGE {0} public-key-file-path string-to-encrypt encrypted-file-path"
    98                .format(sys.argv[0]))
    99          return
   100  
   101      encryptor = Encryptor(sys.argv[1])
   102      data = sys.argv[2]
   103      print(data)
   104      result = encryptor.encrypt(data)
   105      print(result)
   106      with open(sys.argv[3], 'w+') as f:
   107          f.write(data)
   108          f.write('\n')
   109          f.write(result)
   110  
   111  
   112  if __name__ == '__main__':
   113      main()