def __init__(self, reset_callback=None): super(PacketForwarder, self).__init__() self._stop = Event() self.packets_received = 0 # float cuz this can get big self.sending_queue = Queue() self.callback_dict = dict() self.reset_callback = reset_callback self.pc = PacketController() self.pipe = self.pc.pipe_outside
class PacketForwarder(Thread): def __init__(self, reset_callback=None): super(PacketForwarder, self).__init__() self._stop = Event() self.packets_received = 0 # float cuz this can get big self.sending_queue = Queue() self.callback_dict = dict() self.reset_callback = reset_callback self.pc = PacketController() self.pipe = self.pc.pipe_outside def stop(self): self._stop.set() def empty_queue(self): while not self.sending_queue.empty(): try: self.sending_queue.get(False) except Empty: pass def forward_requests(self): try: packet, callback = self.sending_queue.get(block=False) self.callback_dict[packet[:2]] = callback self.pipe.send(packet) self.sending_queue.task_done() except Empty: pass def callback_responses(self): while self.pipe.poll(): request, response = self.pipe.recv() self.packets_received += 1 if response == c.host.reset_msg: self.empty_queue() if self.reset_callback: self.reset_callback() elif len(response) == 1 and response in c.responses: if response == 'G': continue print("Firmware responded with an error: {} " "for the request {}").format(c.responses[response].msg, request) elif request in self.callback_dict: self.callback_dict[request](request, response) def run(self): self.pc.start() while not self._stop.isSet(): self.forward_requests() self.callback_responses() # Kill the pc process if thread stopped self.pc.stop()
from scapy.all import * from scapy.layers.inet import IP, TCP, UDP from datetime import datetime from basic_packet_info import BasicPacketInfo from basic_flow import BasicFlow from flow_generator import FlowGenerator from packet_controller import PacketController from ml_controller import MachineLearningController from notify_controller import NotifyController from predictor.predictor_factory import PredictorFactory import signal import sys pkt_controller = PacketController() ml_controller = MachineLearningController() notify_controller = NotifyController() def on_flow_generated(flow=BasicFlow()): ml_controller.put((flow.flow_id, flow.dump_features())) del [flow] def on_notify(msg): notify_controller.send(msg) def on_pkt(pkt): global pkt_controller pkt_controller.put(pkt) del pkt