SSH-Master-Thesis/ethan_data_processing_scripts/Vector.py

#!/usr/bin/python3
try:
    import cPickle as pickle
except ImportError:
    import pickle
import os
import sys
import sample
import typing
from typing import List
sys.path.insert(0, os.path.dirname(os.path.realpath(__file__)) + \
                '/../src/feature-extractor')


class FeatureVector:
    def __init__(self, s: sample):
        self.activefeatures = []
        # list of key, value tuples represnting values for features.
        self.sampleInfo = s
        if isinstance(s, sample.Sample):
            self.classification = s.user
        else:
            self.classification = "DUMMY DATA - DO NOT USE"

    # set which features are active using a binary list
    # ~~ex: For 2nd, 4th and 5th features to be active, pass [0,1,0,1,1]~~
    # ex: ["total_time", "average_iat", "dead_time"]
    def set_features(self, features: typing.List[str]):
        self.activefeatures = [self.sampleInfo[feature] for feature in features]

    def __repr__(self):
        return str(self.activefeatures)


def writePickledData(filename):
    v = FeatureVector()
    v.features = [1, 2, 3, 4]
    v.activefeatures = [1, 2, 3, 4]
    vs = [v,v]
    with open(filename, 'wb') as file:
        pickle.dump(vs, file)


def readPickledData(filename):
    with open(filename, 'rb') as file:
        x = pickle.load(file)
        # print(x)
    return x


def main():
    # fv = FeatureVector()
    # writePickledData("test.bin")
    # readPickledData("test.txt")
    # s = sample.Sample([{"delta": 0, "time": '09/19/18 13:55:26', }], open("./results.txt"))
    # s.dead_time = 30
    # s.average_iat = 3
    # s.total_time = 30
    s = {"dead_time": 30, "average_iat": 3, "total_time": 30}
    fv = FeatureVector(s)
    fv.set_features(["total_time", "average_iat", "dead_time"])
    print(fv)

if __name__ == '__main__':
    main()