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 typing
from typing import List
sys.path.insert(0, os.path.dirname(os.path.realpath(__file__)) + \
                '/../src/feature-extractor')
import sample

class FeatureVector:
    def __init__(self, s: sample, features = None):
        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"
        if features is not None:
            self.set_features(features)

    # 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]
        return self
    
    def get(self):
        return self.activefeatures

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

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()