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)
class RemoteController(object): def __init__(self, int_listening_port): # Socket properties self.socket = None self.active = True self.connection = None self.client_address = None self.data = None self.buffer_size = long(2 ** 30) # Placeholder for scenario folder self.scenario_folder = None # Placeholder for scenario file self.scenario_file_path = None # Placeholder for scenario name self.scenario_name = None # Placeholder for scenario results file path self.results_file_path = None # Placeholder for controller IP address self.controller_ip_address = None # Standard port, this value will be eventually overridden self.controller_port = 7878 # Spawn quota for local spawner self.spawn_quota = None # Spawner object self.spawner = None # Display options self.stdout_title_width = 64 # Initialize the socket self.__initialize(int_listening_port) return 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 __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 __start_scenario(self): self.spawner.start() click.secho(utils.logify('Launching Spawner...'), fg='green') 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 shutdown(self): self.active = False def __mark_as_title(self, str_message): str_formatted_message = (' '.join(('', str_message, ''))).center(self.stdout_title_width, '=') return ''.join(('\n', str_formatted_message))