Brute Variational Autoencoder

In this approach, we use a VAE to generate entire binaries.
This 'brute' approach is an experiment to see if it can generate functional binaries. Although unlikely to work, it will be interesting to see how far we can get before worrying about feature extraction or metadata interpolation (e.g., extract PE headers and correct the metadata of a generated binary).
Code
Import and preprocess
Here we use the ELFs generated from our earlier work and normalize to 300 byte lengths, using \x90 NOPs as filler.
import numpy as np
def preprocess_samples(samples):
max_length = 300
processed_samples = []
for sample in samples:
if len(sample) < max_length:
sample += b'\x90' * (max_length - len(sample))
processed_samples.append(np.array(list(sample), dtype=np.float32) / 255.0)
return np.array(processed_samples)
import os
def load_binary_files(directory):
samples = []
for filename in os.listdir(directory):
filepath = os.path.join(directory, filename)
if os.path.isfile(filepath):
with open(filepath, 'rb') as file:
binary_data = file.read()
samples.append(binary_data)
return samples
directory = 'aimwg-ph/'
samples = load_binary_files(directory)
print(samples[:5])
print(len(samples))
pp = preprocess_samples(samples)
Compile the model
from tensorflow.keras import layers, models, backend as K
def sampling(args):
z_mean, z_log_var = args
batch = K.shape(z_mean)[0]
dim = K.int_shape(z_mean)[1]
epsilon = K.random_normal(shape=(batch, dim))
return z_mean + K.exp(0.5 * z_log_var) * epsilon
input_dim = 300
intermediate_dim = 64
latent_dim = 2
inputs = layers.Input(shape=(input_dim,))
x = layers.Dense(intermediate_dim, activation='relu')(inputs)
z_mean = layers.Dense(latent_dim)(x)
z_log_var = layers.Dense(latent_dim)(x)
z = layers.Lambda(sampling)([z_mean, z_log_var])
latent_inputs = layers.Input(shape=(latent_dim,))
x = layers.Dense(intermediate_dim, activation='relu')(latent_inputs)
outputs = layers.Dense(input_dim, activation='sigmoid')(x)
encoder = models.Model(inputs, [z_mean, z_log_var, z], name='encoder')
decoder = models.Model(latent_inputs, outputs, name='decoder')
outputs = decoder(encoder(inputs)[2])
vae = models.Model(inputs, outputs, name='vae')
reconstruction_loss = K.mean(K.binary_crossentropy(inputs, outputs)) * input_dim
kl_loss = -0.5 * K.sum(1 + z_log_var - K.square(z_mean) - K.exp(z_log_var), axis=-1)
vae_loss = K.mean(reconstruction_loss + kl_loss)
vae.add_loss(vae_loss)
vae.compile(optimizer='adam')
Train and evaluate
from sklearn.model_selection import train_test_split
X_train, X_test = train_test_split(pp, test_size=0.2, random_state=42)
print("Training shape:", X_train.shape)
print("Testing shape:", X_test.shape)
vae.fit(X_train, X_train,epochs=500,batch_size=32, validation_data=(X_test, X_test))
loss = vae.evaluate(X_test, X_test, batch_size=32)
print("Reconstruction loss:", loss)
Generate new samples
import numpy as np
def sample_latent_points(latent_dim, num_samples):
return np.random.normal(loc=0.0, scale=1.0, size=(num_samples, latent_dim))
def generate_samples(decoder, latent_points):
generated_data = decoder.predict(latent_points)
return generated_data
latent_dim = 2
num_samples = 10
latent_points = sample_latent_points(latent_dim, num_samples)
generated_samples = generate_samples(decoder, latent_points)
def postprocess_binary_samples(samples):
samples = np.round(samples * 255).astype(np.uint8)
return samples
generated_binaries = postprocess_binary_samples(generated_samples)
!mkdir generated/
def save_generated_binaries(generated_binaries, output_dir):
for i, sample in enumerate(generated_binaries):
filepath = os.path.join(output_dir, f"generated_binary_{i}.bin")
with open(filepath, 'wb') as file:
file.write(sample)
output_dir = 'generated/'
save_generated_binaries(generated_binaries, output_dir)
Try and use them!
!ls -la generated/
!file generated/*
total 48
drwxr-xr-x 2 root root 4096 Mar 9 01:47 .
drwxr-xr-x 1 root root 4096 Mar 9 02:02 ..
-rw-r--r-- 1 root root 300 Mar 9 02:27 generated_binary_0.bin
-rw-r--r-- 1 root root 300 Mar 9 02:27 generated_binary_1.bin
-rw-r--r-- 1 root root 300 Mar 9 02:27 generated_binary_2.bin
-rw-r--r-- 1 root root 300 Mar 9 02:27 generated_binary_3.bin
-rw-r--r-- 1 root root 300 Mar 9 02:27 generated_binary_4.bin
-rw-r--r-- 1 root root 300 Mar 9 02:27 generated_binary_5.bin
-rw-r--r-- 1 root root 300 Mar 9 02:27 generated_binary_6.bin
-rw-r--r-- 1 root root 300 Mar 9 02:27 generated_binary_7.bin
-rw-r--r-- 1 root root 300 Mar 9 02:27 generated_binary_8.bin
-rw-r--r-- 1 root root 300 Mar 9 02:27 generated_binary_9.bin
generated/generated_binary_0.bin: ELF 32-bit LSB executable, Intel 80386, version 1 (SYSV), statically linked, no section header
generated/generated_binary_1.bin: ELF 32-bit LSB executable, Intel 80386, version 1 (SYSV), statically linked, no section header
generated/generated_binary_2.bin: ELF 32-bit LSB executable, Intel 80386, version 1 (SYSV), statically linked, no section header
generated/generated_binary_3.bin: data
generated/generated_binary_4.bin: ELF 32-bit LSB executable, Intel 80386, version 1 (SYSV), statically linked, no section header
generated/generated_binary_5.bin: ELF 32-bit LSB executable, Intel 80386, version 1 (SYSV), statically linked, no section header
generated/generated_binary_6.bin: ELF 32-bit LSB executable, Intel 80386, version 1 (SYSV), statically linked, no section header
generated/generated_binary_7.bin: ELF 32-bit LSB executable, Intel 80386, version 1 (SYSV), statically linked, no section header
generated/generated_binary_8.bin: ELF 32-bit LSB executable, Intel 80386, version 1 (SYSV), statically linked, no section header
generated/generated_binary_9.bin: ELF 32-bit LSB executable, Intel 80386, version 1 (SYSV), statically linked, no section header
Interestingly, we have created something resembling a binary. The model has learned the first few bytes of the binary, enough to fool the file command, but is lacking section headers.
Predictably, execution of any of these binaries results in immediate failure (although for one sample it actually generates a segfault, which oddly feels like great progress). Debugging is equally unfruitful.
The brute approach is a fun experiment, but is doomed to failure as we haven't addressed any specific features of the binary. I suspect it's possible to create a 'fixer' application that takes this raw unstructured ELF and reformats it into an executable binary, but then what's the point of training a model to do the heavy lifting for us.
Let's move on to GANs!