def is_agents_up(self): """Whether expected agents all started""" # TrafStatChk.lgr.info("Check for agents availability") try: expected = self.expected_agents available = self.available_agents if len(self.node2aid) == 0: return False if len(expected) != len(available): return False avail = [ai.host for ai in available.values()] notup = [a for a in expected if a not in avail] if len(notup) != 0: to_console("Waiting for {}", notup) return False # for ai in available.values(): # if ai.host not in expected: # return False return True except Exception as ex: TrafStatChk.lgr.error( to_string("Failed to check agent availability", ex)) return False
def on_topic_status(client, userdata, message): """Status topic handler""" self = userdata self.lgr.info(to_string("On topic status")) data = extract_payload(message.payload) if data.get(T_TRAFFIC_STATUS) == TrafficStatus.REGISTERED.name: self.lgr.info(to_string("Traffic {} registered on {}",\ data.get(T_TID), data.get(T_CLIENTID))) self.traffic_state.update_stat(data.get(T_TID), TrafficStatus.REGISTERED, "is_traffic_registered") elif data.get(T_TRAFFIC_STATUS) == TrafficStatus.RUNNING.name: self.lgr.info(to_string("Traffic {} is running on {}",\ data.get(T_TID), data.get(T_CLIENTID))) self.traffic_state.update_stat(data.get(T_TID), TrafficStatus.RUNNING) elif data.get(T_TRAFFIC_STATUS) == TrafficStatus.FINISHED.name: self.lgr.info(to_string("Traffic {} finished", data.get(T_TID))) self.traffic_state.update_stat(data.get(T_TID), TrafficStatus.FINISHED, "is_phase_finished") elif data.get(T_TRAFFIC_STATUS) == TrafficStatus.INITFAILED.name: self.lgr.error(to_string("Agent {} failed to initialize: {}",\ data.get(T_CLIENTID), data.get(T_MSG))) self.teardown() elif data.get(T_TRAFFIC_STATUS) == TrafficStatus.AGENTFAILED.name: self.lgr.error(to_string("Failure on Agent: {}", data)) to_console("Agent failure: {}", data) self.teardown() elif data.get(T_MSG) is not None: self.lgr.info(data.get(T_MSG)) else: self.lgr.debug(to_string("Unhandled status message {}", data))
def start(self): """Start controller routine""" if not ManualCtrl.initialized: self.lgr.error(to_string("ManualCtrl not initialized")) return try: self.lgr.info(to_string("Expected Agents({}): {}",\ len(self.expected_agents), self.expected_agents)) if self.worker is not None: self.worker.start() self.alive_notification() ret = mqtt.MQTT_ERR_SUCCESS for retry_cnt in range(1, 4): ret = self.mqclient.loop_start() if ret is None or ret == mqtt.MQTT_ERR_SUCCESS: break self.mqclient.loop_stop() self.lgr.warning(to_string('Retry [{}] ...', retry_cnt)) if ret is not None and ret != mqtt.MQTT_ERR_SUCCESS: raise RuntimeError("Start MQTT loop failed") self.term.start() to_console("\nBye~\n") self.mqclient.loop_stop() except KeyboardInterrupt: self.lgr.info(to_string("ManualCtrl received keyboard interrupt")) self.term.on_cmd_quit() except NameError as ex: self.lgr.error(to_string("ERR: {}", ex)) except Exception as ex: self.lgr.error(to_string("Exception on ManualCtrl: {}", ex)) self.teardown() self.lgr.info(to_string("ManualCtrl terminated"))
def do_start(self): """Start traffic deployment""" if not self.traf_running: return TrafficState.lgr.info( to_string("Start traffic sm: {}", str(self.current_state))) try: self.create() TrafficState.lgr.info("Starting agents ...") to_console("Starting agents ...") start_agents(self.chk.expected_agents) self.wait_agent() try: while self.traf_running: self.current_phase = next(self.next_phase) gid = self.record_phase() if gid is not None: start_sniffer() self.start_phase() except StopIteration: TrafficState.lgr.info("All traffic generator done") # [self.start_phase() for tg in self.next_phase] except KeyboardInterrupt: TrafficState.lgr.info("Terminating ...") except Exception as ex: TrafficState.lgr.error(to_string("Traffic runtime error: {}", ex)) self.finish()
def run(self, *args, **kwargs): """Execute handler""" if self.handler is not None: try: self.handler(*args, **kwargs) except Exception as ex: to_console(TSH_ERR_HANDLING_CMD, self.name, ex)
def handle_args(): """Arguments handler""" args = None try: rt.parser.description = 'BLE-Zigbee Combo Report Generator' rt.parser.add_argument('-gid', '--genid', required=True, help='Traffic generation identifier') rt.parser.add_argument('-n', '--nodes', help='Nodes of interest', nargs="+") rt.parser.add_argument('-gw', '--gateways', help='Gateway node', nargs="+") rt.parser.add_argument('-csv', '--csv_path', help='Path to CSV data source') args = rt.handle_args() to_console("GENID {}", args.genid) to_console("Nodes of interest {}", args.nodes) to_console("Subnet gateway {}", args.gateways) to_console("CSV Path {}", args.csv_path) except Exception as ex: to_console("Exception on handlin report arguments: {}", ex) return None return args
def init(host, port, sniconf, pcap_file=None, genid='na'): """ if @pcap_file is not `None`, read data from @pcap_file """ logging.info("Pcap2MQ initialization") try: if sniconf.get(T_TSHARK) is None: return False PcapToMq.genid = genid to_console("GENID: {}", PcapToMq.genid) PcapToMq.__create_mqtt_client(host, port) cmd = [] if pcap_file is None: cmd = [ sniconf.get(T_TSHARK), '-l', '-T', 'json', '-i', Sniffer.fifo_path ] else: cmd = [ sniconf.get(T_TSHARK), '-l', '-T', 'json', '-r', pcap_file ] PcapToMq.proc = sp.Popen(cmd, stdout=sp.PIPE, stderr=sp.STDOUT, shell=False, universal_newlines=True) except Exception as ex: logging.error(to_string("Pcap to MQ init failed: {}", ex)) return False return True
def load_storage(self): self.lgr.info(to_string("Load DB plugin:{}", rt.tcfg.storagetype)) try: self._storage = Storage(rt.tcfg, self.lgr) except Exception as ex: to_console("Failed to load plugin: {}", ex) self.lgr.error(to_string("Failed to load plugin: {}", ex))
def on_cmd_traffic_run(self, *args, **kwargs): """Run traffic command handler""" self.lgr.info("On command traffic run") to_console("Start firing") self.ctrl.traffic_state.fire() self.ctrl.traffic_state.phase_finish() to_console("One traffic fired")
def on_cmd_traffic_register(self, *args, **kwargs): """Register traffic command handler""" self.lgr.info("On command traffic register") to_console("Registering traffic") self.ctrl.traffic_state.register() self.ctrl.traffic_state.wait_traffic() to_console("Traffic registerd")
def on_topic_reply(client, userdata, message): """Heartbeat message format: """ self = userdata self.lgr.info(to_string("On topic reply")) data = extract_payload(message.payload) to_console("{}:\n{}", data.get(T_NODE), data.get(T_MSG))
def handle_args(): """Arguments handler""" args = None try: rt.parser.description = 'Report Adapter' rt.parser.add_argument('-proto', required=True, help=\ 'Generate report based on protocol', choices=SupportedProto.protocols.keys()) rt.parser.add_argument('-gid', '--genid', required=True, help='Traffic generation identifier') rt.parser.add_argument('-n', '--nodes', help='Nodes of interest', nargs="+") rt.parser.add_argument('-gw', '--gateways', help='Gateway node', nargs="+") rt.parser.add_argument('-csv', '--csv_path', help='Path to CSV data source') args = rt.handle_args() to_console("GENID {}", args.genid) to_console("Nodes of interest {}", args.nodes) to_console("Subnet gateway {}", args.gateways) to_console("CSV Path {}", args.csv_path) except Exception as ex: to_console("Exception on handlin report arguments: {}", ex) return None return args
def start(self): """Start inuithy shell""" self.lgr.info("Start console") to_console(self._title) to_console(self._banner) mod = os.path.exists(rt.tcfg.tsh_hist) and 'a+' or 'w+' with open(rt.tcfg.tsh_hist, mod) as tshhist: self.console_loop(tshhist)
def do_create(self): if not self.traf_running: return TrafficState.lgr.info( to_string("Create traffic from configure: {}", str(self.current_state))) to_console("Loading traffics") self.phases = create_phases(rt.trcfg, rt.nwcfg) TrafficState.lgr.info(to_string("Total phase: [{}]", len(self.phases))) self.next_phase = self.yield_traffic()
def on_cmd_traffic(self, *args, **kwargs): """Traffic command handler""" self.lgr.info("On command traffic") if args is None or len(args) == 0: to_console(str(self.usages['usage_traffic'])) return params = args[1:] if self._cmd_traffic_routes.get(args[0]): self._cmd_traffic_routes[args[0]](*(params)) else: to_console(str(self.usages['usage_traffic']))
def _setup_banner(self): try: self._banner_path = os.path.abspath(INUITHY_ROOT+'/banner/*') banners = glob.glob(self._banner_path) if banners is None or len(banners) == 0: return self._banner_path = banners[randint(0, len(banners)-1)] with open(self._banner_path, 'r') as fd: self._banner = fd.read() except Exception as ex: to_console("Setup banner failed: {}", ex)
def on_cmd_quit(self, *args, **kwargs): """Quit command handler""" self.lgr.info("On command quit") try: self.running = False if self.ctrl and self.ctrl.traffic_state: # self.ctrl.traffic_state.finish() self.ctrl.traffic_state.phases.clear() self.ctrl.teardown() except Exception as ex: to_console("Exception on quit: {}", ex)
def on_cmd_agent(self, *args, **kwargs): """Operation on agent""" self.lgr.info("On command agent") if args is None or len(args) == 0: to_console(str(self.usages['usage_agent'])) return params = args[1:] if self._cmd_agent_routes.get(args[0]): self._cmd_agent_routes[args[0]](*(params)) else: to_console(str(self.usages['usage_agent']))
def on_cmd_traffic_deploy(self, *args, **kwargs): """Traffic deploy command handler self.deploy, self.wait_nwlayout, self.register, self.wait_traffic, self.fire, self.phase_finish, self.genreport, """ self.lgr.info("On command traffic deploy") # self.ctrl.traffic_state.next() self.ctrl.traffic_state.chk.clear_all() to_console("Start deploying network layout") self.ctrl.traffic_state.deploy() self.ctrl.traffic_state.wait_nwlayout() to_console("Netowrk deployment finished")
def do_waitfor_agent_all_up(self): """Wait for expected agents startup""" if not self.traf_running: return TrafficState.lgr.info( to_string("Wait for agents get ready", str(self.current_state))) to_console("Wait for agents get ready") try: self.chk._is_agents_up.wait() except KeyboardInterrupt: TrafficState.lgr.info("Terminating ...") except Exception as rex: TrafficState.lgr.error(to_string("Wait for agent up: {}", ex))
def on_cmd_agent_start(self, *args, **kwargs): """Agent start command handler""" self.lgr.info("On command agent start") if args is None or len(args) < 1: return hosts = list(args) if T_EVERYONE in hosts: agents = copy.deepcopy(self.ctrl.expected_agents) else: agents = hosts #TODO check for non-configured hosts start_agents(agents) to_console("Waiting for agents to get ready") self.ctrl.traffic_state.wait_agent()
def record_phase(self): """Record running phase""" if not self.traf_running: return nwlayoutname = getnwlayoutname(self.current_phase.nwlayoutid) self.current_phase.genid = str(int(time.time())) cfg = { T_GENID: self.current_phase.genid, T_NWLAYOUT: deepcopy(rt.nwcfg.config.get(nwlayoutname)), } self.ctrl.storage.insert_config(cfg) to_console("Current phase: {}", self.current_phase) self.record_genid(self.current_phase.genid) return self.current_phase
def do_finish(self): """All traffic finished""" TrafficState.lgr.info( to_string("All finished: {}", str(self.current_state))) # if not self.traf_running: # return try: self.phases.clear() TrafficState.lgr.info("Stopping agents ...") stop_agents(self.ctrl.mqclient, rt.tcfg.mqtt_qos) self.chk.set_all() if len(self.chk.available_agents) > 0: TrafficState.lgr.info("Wait for last notifications") self.chk._is_agents_unregistered.wait( int(rt.trcfg.config.get(T_TRAFFIC_FINISH_DELAY))) except KeyboardInterrupt: TrafficState.lgr.info("Terminating ...") to_console("Terminating ...") # self.chk.set_all() except Exception as ex: TrafficState.lgr.error( to_string("Exception on stopping agents: {}", ex)) try: to_console("Stopping sniffer ...") stop_sniffer() except Exception as ex: TrafficState.lgr.error( to_string("Exception on stopping sniffer: {}", ex)) TrafficState.lgr.info("Stopping controller ...") self.ctrl.teardown() # self.chk.done.wait(rt.trcfg.config.get(T_TRAFFIC_FINISH_DELAY)) # self.chk.done.clear() try: if len(self.chk.available_agents) > 0: TrafficState.lgr.info("Force stopping agents ...") force_stop_agents(self.chk.expected_agents) except Exception as ex: TrafficState.lgr.error( to_string("Exception on force stopping agents: {}", ex)) TrafficState.lgr.info("Agents stopped")
def dispatch(self, command): self.lgr.info("Dispatch command") if command is None or len(command) == 0: return cmds = valid_cmds(command) if len(cmds) == 0 or len(cmds[0]) == 0: to_console(TSH_ERR_INVALID_CMD, command, 'help') params = cmds[1:] try: if self._cmd_routes.get(cmds[0]): self._cmd_routes[cmds[0]](*(params)) else: self.on_cmd_help() except Exception as ex: to_console(TSH_ERR_HANDLING_CMD, command, ex)
def handle_args(in_args=None): """Arguments handler""" args = None try: rt.parser.description = 'Inuithy Sniffer' rt.parser.add_argument('-gid', '--genid', required=True, help='Traffic generation identifier') args = rt.handle_args() load_configs() except Exception as ex: to_console("Exception on handling sniffer arguments: {}", ex) return None return args
def do_waitfor_traffic_set(self): """Wait for traffic all registered on agents""" if not self.traf_running: return TrafficState.lgr.info( to_string("Wait for traffic registered: {}", str(self.current_state))) to_console("Wait for traffic registered") try: self.chk._is_traffic_registered.wait() except KeyboardInterrupt: TrafficState.lgr.info("Terminating ...") except Exception as rex: TrafficState.lgr.error( to_string("Wait for get traffic registered: {}", ex))
def do_waitfor_nwlayout_done(self): """Wait for network configure to be done""" if not self.traf_running: return TrafficState.lgr.info( to_string("Wait for network layout done: {}", str(self.current_state))) to_console("Wait for network layout done") try: self.chk._is_nwlayout_done.wait() except KeyboardInterrupt: TrafficState.lgr.info("Terminating ...") except Exception as ex: TrafficState.lgr.error( to_string("Wait for network layout done: {}", ex))
def handle_args(in_args = None): """Arguments handler""" args = None try: rt.parser.description = 'Inuithy Controller' rt.parser.add_argument('-m', '--work-mode', help='Mode for framework to run in', choices=WorkMode.__members__.keys()) args = rt.handle_args() load_tcfg(rt.tcfg_path) if args.work_mode is not None: rt.tcfg.workmode = args.work_mode to_console("Starting {} Controller", rt.tcfg.workmode) except Exception as ex: to_console("Exception on handlin report arguments: {}", ex) return None return args
def do_register(self): """Register traffic to agents Ex: phase0 [tg0, tg1, tg2] phase1 [tg5, tg3] """ if not self.traf_running: return TrafficState.lgr.info( to_string("Register traffic task: {}", str(self.current_state))) to_console("Registering traffic") phase = self.current_phase TrafficState.lgr.info(to_string("Register traffic: [{}]", str(phase))) if not publish_phase(phase, self.chk, self.ctrl.mqclient, rt.tcfg.enable_localdebug): self.traf_running = False
def do_fire(self): """Tell agents to fire registerd traffic""" if not self.traf_running: return TrafficState.lgr.info( to_string("Fire traffic: {}", str(self.current_state))) to_console("Firing traffics ...") for agent in self.chk.available_agents.keys(): if not self.traf_running: break to_console(to_string("Fire on {}", agent)) data = { T_CLIENTID: agent, T_TRAFFIC_TYPE: TrafficType.START.name, } pub_traffic(self.ctrl.mqclient, rt.tcfg.mqtt_qos, data)