vrr/nfstream
2
0
Fork 0
mirror of https://github.com/nfstream/nfstream.git synced 2026-05-19 16:28:14 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

do not return on update, rename raw_packets_matrix to bidirectional_packets_matrix

Browse source
This commit is contained in:
Radion Bikmukhamedov 2020-07-26 19:50:36 +03:00
parent 8a97cb9049
commit 162baadc52
2 changed files with 23 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

10
nfstream/plugin.py
View file

@ -1294,7 +1294,7 @@ class dst2src_fin_packets(NFPlugin):
entry.dst2src_fin_packets += 1
class raw_packets_matrix(NFPlugin):
class bidirectional_packets_matrix(NFPlugin):
"""
WARNING: do not use with .to_pandas() method, rather process the matrix
in the loop, since it will break the DataFrame's row structure.
@ -1398,13 +1398,13 @@ class raw_packets_matrix(NFPlugin):
return self._fill_flow_stats(obs, raw_feature_matrix)
def on_update(self, obs, entry):
if entry.bidirectional_packets > self.packet_limit:
return entry.raw_packets_matrix
return self._fill_flow_stats(obs, entry.raw_packets_matrix, counter=entry.bidirectional_packets - 1)
if entry.bidirectional_packets <= self.packet_limit:
self._fill_flow_stats(obs, entry.bidirectional_packets_matrix, counter=entry.bidirectional_packets - 1)
def on_expire(self, entry):
# rm unfilled matrix rows
entry.raw_packets_matrix = entry.raw_packets_matrix[entry.raw_packets_matrix[:, self.ts_idx] > 0]
unfilled_indexer = entry.bidirectional_packets_matrix[:, self.ts_idx] > 0
entry.bidirectional_packets_matrix = entry.bidirectional_packets_matrix[unfilled_indexer]
"""--------------------------------- nfstream core plugins ----------------------------------------------------------"""

35
tests.py
View file

@ -24,7 +24,8 @@ import csv
import numpy as np
from nfstream import NFStreamer, NFPlugin
from nfstream.plugin import raw_packets_matrix
from nfstream.plugin import bidirectional_packets_matrix
def get_files_list(path):
files = []
@ -328,39 +329,39 @@ messenger.com')
source='tests/pcap/skype.pcap',
idle_timeout=60,
active_timeout=60,
plugins=[raw_packets_matrix(packet_limit=5)],
plugins=[bidirectional_packets_matrix(packet_limit=5)],
statistics=False
)
for entry in streamer:
assert isinstance(entry.raw_packets_matrix, np.ndarray)
assert entry.raw_packets_matrix.shape[1] == 6
assert isinstance(entry.bidirectional_packets_matrix, np.ndarray)
assert entry.bidirectional_packets_matrix.shape[1] == 6
def test_raw_feature_parsing_customized(self):
streamer = NFStreamer(
source='tests/pcap/skype.pcap',
idle_timeout=60,
active_timeout=60,
plugins=[raw_packets_matrix(packet_limit=5,
payload_len=False,
tcp_flag=False,
ip_proto=False,
custom_extractors=[
lambda x: 1,
lambda x: x.direction,
])],
plugins=[bidirectional_packets_matrix(packet_limit=5,
payload_len=False,
tcp_flag=False,
ip_proto=False,
custom_extractors=[
lambda x: 1,
lambda x: x.direction,
])],
statistics=False
)
for entry in streamer:
assert isinstance(entry.raw_packets_matrix, np.ndarray)
assert isinstance(entry.bidirectional_packets_matrix, np.ndarray)
# we have 3 mandatory + 2 custom features
assert entry.raw_packets_matrix.shape[1] == 5
assert entry.bidirectional_packets_matrix.shape[1] == 5
# this is our constant function
assert entry.raw_packets_matrix[0, 3] == 1
assert entry.bidirectional_packets_matrix[0, 3] == 1
# first observation's direction is always 0
assert entry.raw_packets_matrix[0, 4] == 0
assert entry.bidirectional_packets_matrix[0, 4] == 0
if __name__ == '__main__':
unittest.main()