#!/usr/bin/python

import os, csv
import tensorflow as tf
import numpy as np
import pandas as pd
import helpers

# fix random seed for reproducibility
np.random.seed(7)


#-------------------------- Constants --------------------------#
FLAGS = tf.flags.FLAGS
tf.flags.DEFINE_string(
  "input_dir", os.path.abspath("../data/real_logs"),
  "Input directory containing original JSON data files (default = '../data')"
)
tf.flags.DEFINE_string(
  "output_dir", os.path.abspath("../data"),
  "Output directory for TFrEcord files (default = '../data')")

tf.flags.DEFINE_integer("max_vector_len", 16, "Maximum vector length")

#----------------------------------------------------------------#
TRAIN_PATH = os.path.join(FLAGS.input_dir, "20170618_Belma.log")
TEST_PATH = os.path.join(FLAGS.input_dir, "user1_unauthorized.log")

CURRENT_PATH = TEST_PATH
OUTPUT_FILE = "user1_test_C.csv"
#----------------------------------------------------------------#

### START VOCABULARY FUNCTIONS ###
def tokenizer_fn(iterator):
  return (x.split(" ") for x in iterator)


def create_vocabulary(train_path, test_path):
    print("Creating vocabulary...")

    iter_generator = helpers.create_iter_generator(train_path)
    input_iter = []
    for x in iter_generator:
        input = get_features(x)
        input = " ".join(input)
        input_iter.append(input)

    if (test_path):
        iter_generator = helpers.create_iter_generator(test_path)
        for x in iter_generator:
            input = get_features(x)
            for x in input:
                input_iter.append(x)

    vocab_processor = tf.contrib.learn.preprocessing.VocabularyProcessor(
        FLAGS.max_vector_len,
        tokenizer_fn=tokenizer_fn)
    vocab_processor.fit(input_iter)

    print("Done creating vocabulary.")
    return vocab_processor


def write_vocabulary(vocabulary_processor, outfile):
  with open(outfile, "w") as vocabfile:
    for id in range(len(vocabulary_processor.vocabulary_)):
      word =  vocabulary_processor.vocabulary_._reverse_mapping[id]
      vocabfile.write(word + "\n")
  print("Saved vocabulary to {}".format(outfile))


def create_and_save_vocabulary(train, test="", vocabularyfile="vocabulary.txt", processorfile="vocab_processor.bin"):
    vocabulary = create_vocabulary(train, test)

    # Create vocabulary.txt file
    write_vocabulary(vocabulary, os.path.join(FLAGS.output_dir, vocabularyfile))
    # Save vocab processor
    vocabulary.save(os.path.join(tf.flags.FLAGS.output_dir, processorfile))

    return vocabulary


def restore_vocabulary(filename = os.path.join(tf.flags.FLAGS.output_dir, "vocab_processor.bin")):
    return tf.contrib.learn.preprocessing.VocabularyProcessor.restore(filename)
### END VOCABULARY FUNCTIONS ###


def transform_sentence(sequence, vocab_processor):
    # Maps a single vector input into the integer vocabulary.
    if (type(sequence) is not list):
        sequence = [sequence]

    sequence = [" ".join(sequence)]

    vector = next(vocab_processor.transform(sequence)).tolist()
    vector_len = len(next(vocab_processor._tokenizer(sequence)))
    vector = vector[:vector_len]

    return vector


def get_features(line):
    structure = ["added_or_removed", "hour", "usb_devices", "kernel_modules", "open_sockets", "open_sockets",
                 "processes", "open_files", "logged_in_users", "logged_in_users", "shell_history",
                 "listening_ports", "arp_cache", "arp_cache", "syslog", "syslog"]

    # First feature
    added_or_removed = "2"  # added
    if (line["action"] == "removed"):
        added_or_removed = "1"
    # Second feature
    time = helpers.extract_hour(line["unixTime"])

    # Other osquery features
    columns = line["columns"].values()

    # Compatibility with old shell_history query
    #if (line["name"] == "pack_external_pack_shell_history"):
        #columns = str(helpers._parse_shell_history(columns))

    initial_vector = [added_or_removed, time] + ["0"] * (len(structure) - 2)
    # Put action columns in the right place of vector according to structure
    index = structure.index(line["name"].replace('pack_external_pack_', ''))

    for i in range(len(columns)):
        if (columns[i] == ""):
            initial_vector[index + i] = "0"
        else:
            initial_vector[index + i] = columns[i]

    return initial_vector


"""
    Takes logline in json format and vocabulary object.
    Prepare to extract features from logline.
    Return dictionary containing features vector in key name action
"""
def action_to_vector(line, vocabulary):
    features_vector = get_features(line)

    action = transform_sentence(features_vector, vocabulary)

    return action


def create_csv_file(input_filename, output_filename, vocabulary):
    print("Creating CSV file at {}...".format(output_filename))

    actions = []

    for i, row in enumerate(helpers.create_iter_generator(input_filename)):
        action_transformed = action_to_vector(row, vocabulary)
        actions.append(action_transformed)

    output = pd.DataFrame(data={'action': actions})
    output.to_csv(output_filename, index=False, sep=";", quoting=csv.QUOTE_NONE, quotechar='')

    print("Wrote to {}".format(output_filename))


if __name__ == "__main__":
    if (CURRENT_PATH == TRAIN_PATH):
        vocabulary = create_and_save_vocabulary(TRAIN_PATH, TEST_PATH)
    else:
        vocabulary = restore_vocabulary()

    create_csv_file(
        input_filename=CURRENT_PATH,
        output_filename=os.path.join(tf.flags.FLAGS.output_dir, OUTPUT_FILE),
        vocabulary=vocabulary)

#!/usr/bin/python

import os, csv
import tensorflow as tf
import numpy as np
import pandas as pd
import helpers

# fix random seed for reproducibility
np.random.seed(7)


#-------------------------- Constants --------------------------#
FLAGS = tf.flags.FLAGS
tf.flags.DEFINE_string(
  "input_dir", os.path.abspath("../data/real_logs"),
  "Input directory containing original JSON data files (default = '../data')"
)
tf.flags.DEFINE_string(
  "output_dir", os.path.abspath("../data"),
  "Output directory for TFrEcord files (default = '../data')")

tf.flags.DEFINE_integer("max_vector_len", 16, "Maximum vector length")

#----------------------------------------------------------------#
TRAIN_PATH = os.path.join(FLAGS.input_dir, "20170618_Belma.log")
TEST_PATH = os.path.join(FLAGS.input_dir, "user1_unauthorized.log")

CURRENT_PATH = TEST_PATH
OUTPUT_FILE = "user1_test_C.csv"
#----------------------------------------------------------------#

### START VOCABULARY FUNCTIONS ###
def tokenizer_fn(iterator):
  return (x.split(" ") for x in iterator)


def create_vocabulary(train_path, test_path):
    print("Creating vocabulary...")

    iter_generator = helpers.create_iter_generator(train_path)
    input_iter = []
    for x in iter_generator:
        input = get_features(x)
        input = " ".join(input)
        input_iter.append(input)

    if (test_path):
        iter_generator = helpers.create_iter_generator(test_path)
        for x in iter_generator:
            input = get_features(x)
            for x in input:
                input_iter.append(x)

    vocab_processor = tf.contrib.learn.preprocessing.VocabularyProcessor(
        FLAGS.max_vector_len,
        tokenizer_fn=tokenizer_fn)
    vocab_processor.fit(input_iter)

    print("Done creating vocabulary.")
    return vocab_processor


def write_vocabulary(vocabulary_processor, outfile):
  with open(outfile, "w") as vocabfile:
    for id in range(len(vocabulary_processor.vocabulary_)):
      word =  vocabulary_processor.vocabulary_._reverse_mapping[id]
      vocabfile.write(word + "\n")
  print("Saved vocabulary to {}".format(outfile))


def create_and_save_vocabulary(train, test="", vocabularyfile="vocabulary.txt", processorfile="vocab_processor.bin"):
    vocabulary = create_vocabulary(train, test)

    # Create vocabulary.txt file
    write_vocabulary(vocabulary, os.path.join(FLAGS.output_dir, vocabularyfile))
    # Save vocab processor
    vocabulary.save(os.path.join(tf.flags.FLAGS.output_dir, processorfile))

    return vocabulary


def restore_vocabulary(filename = os.path.join(tf.flags.FLAGS.output_dir, "vocab_processor.bin")):
    return tf.contrib.learn.preprocessing.VocabularyProcessor.restore(filename)
### END VOCABULARY FUNCTIONS ###


def transform_sentence(sequence, vocab_processor):
    # Maps a single vector input into the integer vocabulary.
    if (type(sequence) is not list):
        sequence = [sequence]

    sequence = [" ".join(sequence)]

    vector = next(vocab_processor.transform(sequence)).tolist()
    vector_len = len(next(vocab_processor._tokenizer(sequence)))
    vector = vector[:vector_len]

    return vector


def get_features(line):
    structure = ["added_or_removed", "hour", "usb_devices", "kernel_modules", "open_sockets", "open_sockets",
                 "processes", "open_files", "logged_in_users", "logged_in_users", "shell_history",
                 "listening_ports", "arp_cache", "arp_cache", "syslog", "syslog"]

    # First feature
    added_or_removed = "2"  # added
    if (line["action"] == "removed"):
        added_or_removed = "1"
    # Second feature
    time = helpers.extract_hour(line["unixTime"])

    # Other osquery features
    columns = line["columns"].values()

    # Compatibility with old shell_history query
    #if (line["name"] == "pack_external_pack_shell_history"):
        #columns = str(helpers._parse_shell_history(columns))

    initial_vector = [added_or_removed, time] + ["0"] * (len(structure) - 2)
    # Put action columns in the right place of vector according to structure
    index = structure.index(line["name"].replace('pack_external_pack_', ''))

    for i in range(len(columns)):
        if (columns[i] == ""):
            initial_vector[index + i] = "0"
        else:
            initial_vector[index + i] = columns[i]

    return initial_vector


"""
    Takes logline in json format and vocabulary object.
    Prepare to extract features from logline.
    Return dictionary containing features vector in key name action
"""
def action_to_vector(line, vocabulary):
    features_vector = get_features(line)

    action = transform_sentence(features_vector, vocabulary)

    return action


def create_csv_file(input_filename, output_filename, vocabulary):
    print("Creating CSV file at {}...".format(output_filename))

    actions = []

    for i, row in enumerate(helpers.create_iter_generator(input_filename)):
        action_transformed = action_to_vector(row, vocabulary)
        actions.append(action_transformed)

    output = pd.DataFrame(data={'action': actions})
    output.to_csv(output_filename, index=False, sep=";", quoting=csv.QUOTE_NONE, quotechar='')

    print("Wrote to {}".format(output_filename))


if __name__ == "__main__":
    if (CURRENT_PATH == TRAIN_PATH):
        vocabulary = create_and_save_vocabulary(TRAIN_PATH, TEST_PATH)
    else:
        vocabulary = restore_vocabulary()

    create_csv_file(
        input_filename=CURRENT_PATH,
        output_filename=os.path.join(tf.flags.FLAGS.output_dir, OUTPUT_FILE),
        vocabulary=vocabulary)

#!/usr/bin/python # fix random seed for reproducibility #-------------------------- Constants --------------------------# #----------------------------------------------------------------# #----------------------------------------------------------------# ### START VOCABULARY FUNCTIONS ### # Create vocabulary.txt file # Save vocab processor ### END VOCABULARY FUNCTIONS ### # Maps a single vector input into the integer vocabulary. # First feature # added # Second feature # Other osquery features # Compatibility with old shell_history query #if (line["name"] == "pack_external_pack_shell_history"): #columns = str(helpers._parse_shell_history(columns)) # Put action columns in the right place of vector according to structure Takes logline in json format and vocabulary object.     Prepare to extract features from logline.     Return dictionary containing features vector in key name action