def __setup_scenario(self, dic_payload): # Get the Base64 encoded ZIP file from payload click.secho(utils.logify('Unpacking scenario... '), nl=False) str_base64_zipped_scenario = dic_payload['scenarioBase64ZippedFolder'] obj_zipped_scenario = StringIO(base64.b64decode(str_base64_zipped_scenario)) obj_zipfile = ZipFile(obj_zipped_scenario, 'r') # Extract ZIP file into temporary folder self.scenario_folder = tempfile.mkdtemp(prefix='woodpecker-') obj_zipfile.extractall(self.scenario_folder) click.secho('DONE', fg='green', bold=True) # Logging temp folder str_message = ' '.join(('Scenario temporary folder set to: ', self.scenario_folder)) click.secho(utils.logify(str_message)) # Fill all the object properties self.controller_port = dic_payload.get('controllerPort', '7878') self.spawn_quota = dic_payload.get('spawnQuota', 1) self.scenario_name = dic_payload.get('scenarioName', 'Scenario') self.scenario_file_path = utils.get_abs_path(dic_payload.get('scenarioFile', './scenario.py'), self.scenario_folder) self.results_file_path = utils.get_abs_path(dic_payload.get('resultsFile', './results/results.sqlite'), self.scenario_folder) # Create spawner click.secho(utils.logify('Setting up spawner... ')) self.spawner = Spawner(self.controller_ip_address, self.controller_port, self.scenario_folder, self.scenario_name, self.scenario_file_path, self.results_file_path, self.spawn_quota) click.secho(utils.logify('Spawner created'), fg='green', bold=True)
def serve_forever(self): click.secho(utils.logify('Waiting for controller connection...'), fg='green', bold=True) while self.active: self.connection, self.client_address = self.socket.accept() self.data = self.connection.recv(self.buffer_size) self.__handle() self.connection.close() str_message = ' '.join(('Remote controller gracefully closed by controller on', self.client_address[0])) click.echo(utils.logify(str_message))
def __initialize(self, **kwargs): # Controller and Log collector socket port self.port = kwargs.get('controller_port', 7878) # Controller address, void at the beginning self.server_address = kwargs.get('controller_ip_address', None) # Scenario folder, defaults to current directory self.scenario_folder = kwargs.get('scenario_folder', os.getcwd()) # Scenario file name, defaults to './scenario.py' self.scenario_file_path = utils.get_abs_path(kwargs.get('scenario_file', './scenario.py'), self.scenario_folder) # Results file path, defaults to './results/results.sqlite' self.results_file_path = utils.get_abs_path(kwargs.get('results_file', './results/results.sqlite'), self.scenario_folder) # Scenario name, defaults to 'Scenario' self.scenario_name = kwargs.get('scenario_name', 'Scenario') # Spawn quota for local spawner self.spawn_quota = kwargs.get('spawn_quota', 1) # Placeholder for scenario self.scenario = None # Armed flag, used to start and stop running self.armed = False # Elapsed time since scenario beginning self.elapsed_time = 0 # Internal message sender placeholder self.sender = Sender(self.server_address, self.port, 'UDP') if self.server_address else None # Sender spawn message poll interval self.sender_spawn_polling_interval = kwargs.get('sender_spawn_polling_interval', 5) self.sender_spawn_elapsed_time = 0.0 # Internal Sysmonitor object and related thread self.sysmonitor_polling_interval = kwargs.get('sysmonitor_polling_interval', 1.0) click.secho(utils.logify('Setting up Sysmonitor... ', 'Spawner')) self.sysmonitor = Sysmonitor(self.server_address, self.port, self.sysmonitor_polling_interval) click.secho(utils.logify('Sysmonitor started'), fg='green', bold=True)
def __initialize(self, int_listening_port): # Initialize socket self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) self.listening_port = int_listening_port self.socket.bind(('', self.listening_port)) self.socket.listen(3) # Print message on screen str_message = ' '.join(('Socket opened on port', str(self.listening_port))) click.echo(utils.logify(str_message))
def __handle(self): # Read the messages and parse them into dict obj_message = json.loads(self.data) self.controller_ip_address = self.client_address[0] str_message = ' '.join(('Controller connected from ', self.controller_ip_address)) click.secho(utils.logify(str_message), fg='green') # Switch action according to message type if obj_message['dataType'] == 'setup': click.secho(self.__mark_as_title('SCENARIO SETUP'), fg='red') self.__setup_scenario(obj_message['payload']) elif obj_message['dataType'] == 'start': click.secho(self.__mark_as_title('STARTING SCENARIO'), fg='red') self.__start_scenario() elif obj_message['dataType'] == 'stop': click.secho(self.__mark_as_title('STOPPING SCENARIO'), fg='red') elif obj_message['dataType'] == 'emergency_stop': click.secho(self.__mark_as_title('!!!EMERGENCY STOP!!!'), err=True) elif obj_message['dataType'] == 'shutdown': click.secho(self.__mark_as_title('CLOSING REMOTE CONTROLLER'), fg='red') self.shutdown()
def __start_scenario(self): self.spawner.start() click.secho(utils.logify('Launching Spawner...'), fg='green')
def __run(self): click.secho(utils.logify('Starting Sysmonitor... ', 'Spawner'), nl=False) self.sysmonitor.start() click.secho('DONE', fg='green', bold=True) self.scenario.configure() self.scenario.navigations_definition() self.scenario.scenario_start = utils.get_timestamp(False) self.scenario.scenario_duration =\ self.scenario.get_scenario_duration() list_navigations = self.scenario.get_navigation_names() # Update sender spawn message elapsed time self.sender_spawn_elapsed_time += self.elapsed_time # Write beginning message click.secho(utils.logify('Running scenario...', 'Spawner'), fg='green') while self.elapsed_time <= self.scenario.scenario_duration and self.armed: # Get elapsed time time_elapsed = utils.get_timestamp(False) - self.scenario.scenario_start self.elapsed_time = time_elapsed.total_seconds() # Cycle through navigations for str_navigation_name in list_navigations: # Get planned and current spawn number self.scenario.navigations[str_navigation_name]['planned_spawns'] =\ self.scenario.get_planned_spawns(str_navigation_name, self.elapsed_time) self.scenario.navigations[str_navigation_name]['current_spawns'] =\ len(self.scenario.navigations[str_navigation_name]['threads']) # Send spawn message, but only if is passed enough time from last sending if self.sender_spawn_elapsed_time >= self.sender_spawn_polling_interval: self.sender.send('spawns', { 'hostName': utils.get_ip_address(), 'timestamp': utils.get_timestamp(), 'navigationName': str_navigation_name, 'plannedSpawns': self.scenario.navigations[str_navigation_name]['planned_spawns'], 'runningSpawns': self.scenario.navigations[str_navigation_name]['current_spawns'] }) self.sender_spawn_elapsed_time = 0.0 int_spawns_difference = self.scenario.navigations[str_navigation_name][ 'current_spawns'] - self.scenario.navigations[ str_navigation_name]['planned_spawns'] # If there are less spawns than planned, add some spawns if int_spawns_difference < 0: for int_counter in range(0, -int_spawns_difference): str_navigation_path = self.scenario.get_navigation_path(str_navigation_name) str_id = utils.random_id(16) print(str_id) obj_spawn = Spawn(str_id, str_navigation_name, str_navigation_path, self.scenario_folder, self.server_address, self.port, self.scenario.settings) obj_spawn.start() self.scenario.navigations[str_navigation_name]['threads'].append(obj_spawn) # If there are more spawns than planned, start killing older spawns elif int_spawns_difference > 0: for int_counter in range(0, int_spawns_difference): obj_spawn = self.scenario.navigations[str_navigation_name]['threads'].pop(0) obj_spawn.terminate() obj_spawn.join() print('Killed one') self.__unarm() # Clean all the remaining threads click.secho(utils.logify('Cleaning up... ', 'Spawner'), nl=False) for str_navigation_name in list_navigations: for int_counter in range(0, len(self.scenario.navigations[str_navigation_name]['threads'])): obj_spawn = self.scenario.navigations[str_navigation_name]['threads'].pop(0) obj_spawn.terminate() obj_spawn.join() print('Killed one - Cleanup') click.secho('DONE', fg='green', bold=True) # Close Sysmonitor thread click.secho(utils.logify('Closing Sysmonitor... ', 'Spawner'), nl=False) self.sysmonitor.terminate() self.sysmonitor.join() click.secho('DONE', fg='green', bold=True) click.secho(utils.logify('Scenario gracefully completed!', 'Spawner'), fg='green', bold=True)