0xTriboulet-evading_the_machine/plot_1.py

#!/usr/bin/env python3
"""
plots.py
--------
Reads features from pe_features.csv and plots them in Euclidean space.
Malicious samples (label=1) are colored red.
Benign samples (label=0) are colored blue.
"""

import pandas as pd
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  # needed for 3D projection

def main():
    # Load dataset
    df = pd.read_csv("pe_features.csv")

    # Extract features and labels
    X = df[["f1_entropy", "f2_strings_density", "f3_log_size"]].values
    y = df["label"].values

    # Create 3D scatter plot
    fig = plt.figure(figsize=(10, 7))
    ax = fig.add_subplot(111, projection="3d")

    # Separate benign and malware
    benign = X[y == 0]
    malware = X[y == 1]

    # Plot benign (blue) and malware (red)
    ax.scatter(benign[:, 0], benign[:, 1], benign[:, 2],
               c="blue", label="Benign", alpha=0.6, s=40)
    ax.scatter(malware[:, 0], malware[:, 1], malware[:, 2],
               c="red", label="Malware", alpha=0.6, s=40)

    # Axis labels
    ax.set_xlabel("Entropy")
    ax.set_ylabel("Strings Density")
    ax.set_zlabel("Log(Size)")

    # Title and legend
    ax.set_title("PE Feature Space Visualization")
    ax.legend()

    plt.show()

if __name__ == "__main__":
    main()