ma_cisco_malware/models/networks.py

from collections import namedtuple

import keras
from keras.engine import Input, Model as KerasModel
from keras.layers import Conv1D, Dense, Dropout, Embedding, GlobalMaxPooling1D, TimeDistributed

import dataset

Model = namedtuple("Model", ["in_domains", "in_flows", "out_client", "out_server"])


def get_domain_embedding_model(embedding_size, input_length, filter_size, kernel_size, hidden_dims,
                               drop_out=0.5) -> KerasModel:
    x = y = Input(shape=(input_length,))
    y = Embedding(input_dim=dataset.get_vocab_size(), output_dim=embedding_size)(y)
    y = Conv1D(filter_size,
               kernel_size,
               activation='relu')(y)
    y = GlobalMaxPooling1D()(y)
    y = Dropout(drop_out)(y)
    y = Dense(hidden_dims, activation="relu")(y)
    return KerasModel(x, y)


def get_final_model(cnnDropout, flow_features, window_size, domain_length, cnn_dims, kernel_size,
                    dense_dim, cnn) -> Model:
    ipt_domains = Input(shape=(window_size, domain_length), name="ipt_domains")
    encoded = TimeDistributed(cnn, name="domain_cnn")(ipt_domains)
    ipt_flows = Input(shape=(window_size, flow_features), name="ipt_flows")
    merged = keras.layers.concatenate([encoded, ipt_flows], -1)
    # CNN processing a small slides of flow windows
    y = Conv1D(cnn_dims,
               kernel_size,
               activation='relu')(merged)
    # remove temporal dimension by global max pooling
    y = GlobalMaxPooling1D()(y)
    y = Dropout(cnnDropout)(y)
    y = Dense(dense_dim, activation='relu')(y)
    out_client = Dense(1, activation='sigmoid', name="client")(y)
    out_server = Dense(1, activation='sigmoid', name="server")(y)

    return Model(ipt_domains, ipt_flows, out_client, out_server)


def get_inter_model(dropout, flow_features, window_size, domain_length, cnn_dims, kernel_size,
                    dense_dim, cnn) -> Model:
    ipt_domains = Input(shape=(window_size, domain_length), name="ipt_domains")
    ipt_flows = Input(shape=(window_size, flow_features), name="ipt_flows")
    encoded = TimeDistributed(cnn, name="domain_cnn")(ipt_domains)
    merged = keras.layers.concatenate([encoded, ipt_flows], -1)
    y = Dense(dense_dim,
              activation="relu",
              name="dense_server")(merged)
    out_server = Dense(1, activation="sigmoid", name="server")(y)
    merged = keras.layers.concatenate([merged, y], -1)
    # CNN processing a small slides of flow windows
    y = Conv1D(cnn_dims,
               kernel_size,
               activation='relu')(merged)
    # remove temporal dimension by global max pooling
    y = GlobalMaxPooling1D()(y)
    y = Dropout(dropout)(y)
    y = Dense(dense_dim,
              activation='relu',
              name="dense_client")(y)

    out_client = Dense(1, activation='sigmoid', name="client")(y)

    return Model(ipt_domains, ipt_flows, out_client, out_server)


def get_server_model(flow_features, domain_length, dense_dim, cnn):
    ipt_domains = Input(shape=(domain_length,), name="ipt_domains")
    ipt_flows = Input(shape=(flow_features,), name="ipt_flows")
    encoded = cnn(ipt_domains)
    cnn.name = "domain_cnn"
    
    merged = keras.layers.concatenate([encoded, ipt_flows], -1)
    y = Dense(dense_dim,
              activation="relu",
              name="dense_server")(merged)
    out_server = Dense(1, activation="sigmoid", name="server")(y)
    
    return KerasModel(inputs=[ipt_domains, ipt_flows], outputs=out_server)


def get_long_model(dropout, flow_features, window_size, domain_length, cnn_dims, kernel_size,
                   dense_dim, cnn) -> Model:
    ipt_domains = Input(shape=(window_size, domain_length), name="ipt_domains")
    ipt_flows = Input(shape=(window_size, flow_features), name="ipt_flows")
    encoded = TimeDistributed(cnn, name="domain_cnn")(ipt_domains)
    merged = keras.layers.concatenate([encoded, ipt_flows], -1)
    y = Conv1D(cnn_dims,
               kernel_size,
               activation='relu', name="conv_server")(merged)
    # remove temporal dimension by global max pooling
    y = GlobalMaxPooling1D()(y)
    y = Dropout(dropout)(y)
    y = Dense(dense_dim,
              activation="relu",
              name="dense_server")(y)
    out_server = Dense(1, activation="sigmoid", name="server")(y)
    # CNN processing a small slides of flow windows
    y = Conv1D(cnn_dims,
               kernel_size,
               activation='relu', name="conv_client")(merged)
    # remove temporal dimension by global max pooling
    y = GlobalMaxPooling1D()(y)
    y = Dropout(dropout)(y)
    y = Dense(dense_dim,
              activation='relu',
              name="dense_client")(y)
    
    out_client = Dense(1, activation='sigmoid', name="client")(y)
    
    return Model(ipt_domains, ipt_flows, out_client, out_server)