%env CUDA_VISIBLE_DEVICES=0 # limit GPU usage, if any to this GPU

"""
Explanation: Preliminaries
We're going to build and compare a few malware machine learning models in this series of Jupyter notebooks.  Some of them require a GPU.  I've used a Titan X GPU for this exercise.  If yours isn't as beefy, you may get tensorflow memory errors that may require modifying some of the code, namely file_chunks and file_chunk_size.  (I'll point to it later.)  But, to get started, the first few exercises will work on even that GPU you're embarrassed to tell people about, or if you're willing to wait, no GPU at all.
For the fancy folks who have multiple GPUs, we're going to restrict usage to the first one.
End of explanation
"""


from classifier import common

# this will take a LONG time the first time you run it (and cache features to disk for next time)
# it's also chatty.  Parts of feature extraction require LIEF, and LIEF is quite chatty.
# the output you see below is *after* I've already run feature extraction, so that
#   X and sample_index are being read from cache on disk
X, y, sha256list = common.extract_features_and_persist() 

# split our features, labels and hashes into training and test sets
from sklearn.model_selection import train_test_split
import numpy as np
np.random.seed(123)
X_train, X_test, y_train, y_test, sha256_train, sha256_test = train_test_split( X, y, sha256list, test_size=1000) 
# a random train_test split, but for a malware classifier, we should really be holding out *future* malicious and benign 
# samples, to better capture how we'll generalize to malware yet to be seen in the wild. ...an exercise left to the reader..

"""
Explanation: Also note that this exercise assumes you've already populated a malicious/ and a benign/ directory with samples that you consider malicious and benign, respectively.  How many samples?  In this notebook, I'm using 50K of each for demonstration purposes.  Sadly, you must bring your own.  If you don't populate these subdirectories for binaries (each renamed to the sha256 hash of its contents!), the code will bicker and complain incessently.
Feature extraction for feature-based models
There is a lot of domain knowledge on what malware authors can do, and what malware authors actually do when crafting malicious files.  Furthermore, there are some things malware authors seldom do that would indicate that a file is benign.  For each file we want to analyze, we're going to encapsulate that domain knowledge about malicious and benign files in a single feature vector.  See the source code at classifier/pefeatures.py.
Note that the feature extraction we use here contains many elements from published malware classification papers. Some of those are slightly modified.  And there are additional features in this particular feature extraction that are included because, well, they were just sitting there in the LIEF parser patiently waiting for a chair at the feature vector table.  Read: there's really no secret sauce in there, and to turn this into something commercially viable would take a bit of work.  But, be my guest.
A note about LIEF. What a cool tool with a great mission! It aims to parse and manipulate binary files for Windows (PE), Linux (ELF) and MacOS (macho).  Of course, we're using only the PE subset here.  At the time of this writing, LIEF is still very much a new tool, and I've worked with the authors to help resolve some kinks.  It's a growing project with more warts to find and fix.  Nevertheless, we're using it as the backbone for features that requires one to parse a PE file.
End of explanation
"""


# StandardScaling the data can be important to multilayer perceptron
from sklearn.preprocessing import StandardScaler
scaler = StandardScaler().fit(X_train)

# Note that we're using scaling info form X_train to transform both
X_train = scaler.transform(X_train) # scale for multilayer perceptron
X_test = scaler.transform(X_test)

from classifier import simple_multilayer
from keras.callbacks import LearningRateScheduler, EarlyStopping, ReduceLROnPlateau, ModelCheckpoint
model = simple_multilayer.create_model(
    input_shape=(X_train.shape[1], ),     # input dimensions
    input_dropout=0.05,                   # this prevents the model becoming a fanboy of (overfitting to) any particular input feature
    hidden_dropout=0.1,                   # same, but for hidden units.  Dropping out hidden layers can create a sort of ensemble learner
    hidden_layers=[4096, 2048, 1024, 512] # this is "art". making up # of hidden layers and width of each. don't be afraid to change this
)
model.fit(X_train, y_train,
          batch_size=128,
          epochs=200,
          verbose=1,
          callbacks=[
              EarlyStopping( patience=20 ),
              ModelCheckpoint( 'multilayer.h5', save_best_only=True),
              ReduceLROnPlateau( patience=5, verbose=1)],
          validation_data=(X_test, y_test))

from keras.models import load_model
# we'll load the "best" model (in this case, the penultimate model) that was saved 
# by our ModelCheckPoint callback
model = load_model('multilayer.h5')

y_pred = model.predict(X_test)
common.summarize_performance(y_pred, y_test, "Multilayer perceptron") 
# The astute reader will note we should be doing this on a separate holdout, since we've explicitly
# saved the model that works best on X_test, y_test...an exercise for left for the reader...


"""
Explanation: Multilayer perceptron
We'll use the features we extracted to train a multilayer perceptron (MLP).  An MLP is an artificial neural network with at least one hidden layer.  Is a multilayer perceptron "deep learning"?  Well, it's a matter of semantics, but "deep learning" may imply that the features and model are optimized together, end-to-end.  So, it that sense, no: since we're using domain knowledge to extract features, then pass it to an artificial neural network, we'll remain conservative and call this an MLP. (As we'll see, don't get fooled just because we're not calling this "deep learning": this MLP is no slouch.) The network architecture is defined in classifier/simple_multilayer.py.
End of explanation
"""


from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
# you can increase performance by increasing n_estimators, and removing the restriction on max_depth
# I've kept those in there because I want a quick-and-dirty look at how the MLP above
rf = RandomForestClassifier( 
    n_estimators=40, 
    n_jobs=-1, 
    max_depth=30
).fit(X_train, y_train)

y_pred = rf.predict_proba(X_test)[:,-1] # get probabiltiy of malicious (last class == last column )
_ = common.summarize_performance(y_pred, y_test, "RF Classifier")

"""
Explanation: Sanity check: random forest classifier
Alright.  Is that good?  Let's compare to another model. We'll reach for the simple and reliable random forest classifier?
One nice thing about tree-based classifiers like a random forest classifier is that they are invariant to linear scaling and shifting of the dataset (the model will automatically learn those transformations).  Nevertheless, for a sanity check, we're going to use the scaled/transformed features in a random forest classifier.
End of explanation
"""


