class UIListener: def __init__(self, address, port, authkey): self.client = Client((address, port), authkey=authkey) def get_args(self, n): return [self.client.recv() if self.client.poll(0.5) else None for _ in range(n)] def listenloop(self): keepalive = True try: while keepalive: while self.client.poll(): data = self.client.recv() print('{}: {}'.format('x64' if utils.is64bit() else 'x86', data)) if data == 'close': keepalive = False else: func = _interface[data] self.client.send(func(*self.get_args(func.__code__.co_argcount))) print('{}: sent response to {}'.format('x64' if utils.is64bit() else 'x86', data)) time.sleep(0.05) except EOFError or ConnectionError: pass self.client.close()
def serverrun_single(problem_id, code, input_data, output_data): import acwingcli.config as config url = config.problem_cache[eval(problem_id)]['link'] address = ('localhost', 6001) conn = Client(address, authkey=b'1234') local_server_message = { 'activity': 'run', 'url': url, 'code': code, 'input_data': input_data } conn.send(json.dumps(local_server_message)) while True: early_exit = False if conn.poll(20) == True: response = json.loads(conn.recv()) if 'local_debug_message' in response.keys(): display_debug_message(response) elif 'status' in response.keys(): early_exit = display_run_status(response, problem_id, input_data, output_data) else: sys.stdout.write(Fore.GREEN + Style.BRIGHT + 'TIME OUT' + Style.RESET_ALL) sys.stdout.flush() if early_exit == True: break conn.close()
def shutdown_server(): try: address = ('localhost', 6001) conn = Client(address, authkey=b'1234') local_server_message = {'activity': 'stopserver'} conn.send(json.dumps(local_server_message)) while True: early_exit = False if conn.poll(20) == True: cmdwrite.client_debug( 'shutdown message sent, waiting for response') response = json.loads(conn.recv()) if 'local_debug_message' in response.keys(): display_debug_message(response) if (response['local_debug_message'] == 'close'): early_exit = True cmdwrite.client_debug( 'server shutdown confirmation received') if early_exit == True: break else: cmdwrite.client_debug( 'no server shutdown confirmation received, exit') break except ConnectionRefusedError: cmdwrite.client_debug('no server connection, exit') except EOFError: cmdwrite.client_debug('no server connection, exit')
class VariableMonitoringThread(Thread): def __init__(self, rule): Thread.__init__(self) self.rule = rule self.listener_ready = False def __create_connection(self): tmp_base = utils.get_tmp_base_path_from_mkfile() plc_base_path = os.path.join(tmp_base, self.rule['plc'].name) address = os.path.join(plc_base_path, 'plc_socket') # Wait until listener is ready # TODO: Refactoring needed (anti-pattern) while not self.listener_ready: try: self.conn = Client(address) self.listener_ready = True except socket.error as e: # Check if listener is not yet ready if e.errno == ENOENT: # socket.error: [Errno 2] No such file or directory sleep(1) else: logger.exception("Unknown socket error occurred.") def run(self): self.__create_connection() tag_names = [self.rule['var_name']] self.conn.send(MonitorMessage(tag_names)) while True: try: if self.conn.poll(): result = self.conn.recv() if isinstance(result, UnknownPlcTagException): raise UnknownPlcTagException(result) elif isinstance(result, TerminateMessage): logger.info("Terminating monitoring of '{}'.".format( self.rule['var_name'])) break elif isinstance(result, MonitorResponseMessage): if self.rule['predicate'] == Predicates.MAXVAL: if int(result.value) > self.rule['value']: logger.warning( "ALERT! '{}' tag [{}={}] exceeds max value of {}." .format(self.rule['plc'].name, self.rule['var_name'], result.value, self.rule['value'])) else: logger.error("Received unexpected message type '%s'.", type(result)) except EOFError: logger.exception("Received EOF.") break except UnknownPlcTagException: logger.exception("Unknown PLC Tag.") self.conn.close()
class RemoteAgent: def __init__(self, address, socket_family, connection_retries=100): self._log = logging.getLogger(self.__class__.__name__) self._conn = None self._tp_exec = futures.ThreadPoolExecutor() for i in range(connection_retries): # Waiting on agent to open it's socket. try: self._conn = Client(address, socket_family) break except Exception: self._log.debug( f"RemoteAgent retrying connection to agent in: attempt {i}" ) time.sleep(0.1) if self._conn is None: raise RemoteAgentException("Failed to connect to remote agent") def __del__(self): self.terminate() def _act(self, obs, timeout): # Send observation self._conn.send({"type": "obs", "payload": obs}) # Receive action if self._conn.poll(timeout): try: return self._conn.recv() except ConnectionResetError as e: self.terminate() raise e else: return None def act(self, obs, timeout=None): # Run task asynchronously and return a Future return self._tp_exec.submit(self._act, obs, timeout) def start(self, agent_spec: AgentSpec): # Send the AgentSpec to the agent runner self._conn.send( # We use cloudpickle only for the agent_spec to allow for serialization of lambdas {"type": "agent_spec", "payload": cloudpickle.dumps(agent_spec)} ) def terminate(self): if self._conn: self._conn.close() # Shutdown thread pool executor self._tp_exec.shutdown()
class CliClient(cmd.Cmd): def __init__(self,address,id): self.id = id self.connection = Client(address, authkey='secret password') cmd.Cmd.__init__(self) def onecmd(self,line): self.connection.send(("creator",line)) while self.connection.poll(0.1): print self.connection.recv()
class ProcClient(threading.Thread): def __init__(self, addr, authkey): threading.Thread.__init__(self) self.bus = None self.addr = addr self.authkey = authkey self.conn = None self.running = False self.setDaemon(True) def run(self): self.conn = None for i in range(0, 3): self.bus.log("Starting process client connection to: %s:%d" % self.addr) try: self.conn = Client(self.addr, authkey=self.authkey) break except: self.bus.log("", traceback=True) time.sleep(10.0) if not self.conn: self.bus.log("Failed to connect to %s:%d" % self.addr) return self.running = True while self.running: try: if self.conn.poll(1.0): print self.conn.recv() except IOError: self.stop() break def stop(self): if not self.running: return self.bus.log("Stopping process client connection") if self.conn: self.running = False self.conn.close() self.conn = None def send(self, data): self.conn.send(data) def recv(self): return self.conn.recv()
def controller(args): """ Specifies linear and angular accelerations to be applied by boat. This function is to be executed in the process that handles (or simulates handling) the main control of the boat. """ controller = None if args.controller == "keyboard": controller = KeyboardController(in_sim=False) elif args.controller == "autonomy_template": controller = AutonomyControllerTemplate(in_sim=False) elif args.controller == "complementary_filter_test": controller = ComplementaryFilterController(in_sim=False) elif args.controller == "minimal_controller": controller = MinimalController(in_sim=False) elif args.controller == "scipy_logging": controller = ScipyLoggingController(in_sim=False) elif args.controller == "scipy_opt": controller = ScipyOptController(in_sim=False) elif args.controller == "slsqp": controller = SLSQPController(in_sim=False) elif args.controller == "pid": controller = PIDController(in_sim=False) elif args.controller == "planning": controller = PlanningController(in_sim=False) elif args.controller == "c_planning": controller = ControlPlanner(in_sim=False) address = ('localhost', 8000) conn = Client(address, authkey=b'secret password') radio_address = ('localhost', 8001) radio_conn = Client(radio_address, authkey=b'secret password') waypoints = radio_conn.recv() conn.send(waypoints) env = types.SimpleNamespace(waypoints=waypoints) while True: if conn.poll(): state = json.loads(conn.recv())["state"] state = state["lon"], state["lat"], state["speed"], state["desired_speed"], state["angle"], state["ang_vel"], state["ocean_current_x"], state["ocean_current_y"], state["obstacles"] action = controller.choose_action(env, state) conn.send(action)
class GenericClient2(object): def __init__(self, address, port, key, logger, sname): self.address = address self.port = int(port) self.key = key self.logger = logger self.serverName = sname while True: # keep going until we break out inside the loop try: self.logger.debug('Attempting to connect to ' + self.serverName + ' server at ' + str(self.address) + ' port ' + str(self.port)) self.conn = Client((self.address, self.port)) self.logger.debug('Connect to ' + self.serverName + ' successful.') break except SocketError as serr: if serr.errno == errno.ECONNREFUSED: self.logger.debug( 'Connect to ' + self.serverName + ' failed because connection was refused (the server is down). Trying again.' ) else: # Not a recognized error. Treat as fatal. self.logger.debug('Connect to ' + self.serverName + ' gave socket error ' + str(serr.errno)) raise serr except: self.logger.exception('Connect to ' + self.serverName + ' threw unknown exception') raise def send(self, msg): self.conn.send(json.dumps(msg)) self.logger.debug('send message: ' + str(msg)) def poll(self, t): return self.conn.poll(t) def recv(self): return self.conn.recv() def close(self): self.conn.close()
class ProcessManager(object): def __init__(self, extra_kwargs={}): # Maps entity class name -> (entity class, process name) self.processes = {} # Initialize pipe to simulator self.simulator_pipe = Client(("localhost", 3829)) def start_process(self, entity_class, name, *args, **kwargs): '''Initiates a process of the class proc_cls.''' self.processes[entity_class.__name__] = (entity_class, name) self.simulator_pipe.send( map(lambda x: x[0], self.processes.itervalues()) ) def get_data(self, **kwargs): '''get data from all running processes and package the output into a dictionary ''' vision_data = {} if self.simulator_pipe.poll(): for entity_cls, data in self.simulator_pipe.recv(): entity_cls, process_name = self.processes[entity_cls.__name__] vision_data[process_name] = data if vision_data: for entity_cls_name, process_name in self.processes.itervalues(): if process_name not in vision_data: vision_data[process_name] = None return vision_data else: return None def ping(self): '''verify that all processes are alive''' return True def kill(self): ''' kill all running sub processes ''' self.simulator_pipe.send([]) self.processes = {}
def dispatch(*args): '''If there is an existing zim process pass along the arguments @param args: commandline arguments @raises AssertionError: when no existing zim process or connection failed ''' assert not get_in_main_process() try: logger.debug('Connecting to %s', SERVER_ADDRESS) conn = Client(SERVER_ADDRESS, SERVER_ADDRESS_FAMILY) conn.send(args) if conn.poll(5): re = conn.recv() else: re = 'No response' except Exception, e: if hasattr(e, 'errno') and e.errno == 2: raise AssertionError, 'No such file or directory' else: raise AssertionError, 'Connection failed'
def serversubmit(problem_id, code): total_attempt = 0 while total_attempt <= 5: total_attempt += 1 try: import acwingcli.config as config url = config.problem_cache[eval(problem_id)]['link'] address = ('localhost', 6001) conn = Client(address, authkey=b'1234') local_server_message = { 'activity': 'send', 'url': url, 'code': code } conn.send(json.dumps(local_server_message)) except ConnectionRefusedError: if len(utils.get_acwing_server_process()) == 0: cmdwrite.status('No server found. Initializing server...') subprocess.Popen(['acwingcli', '-runserver'], stdout=subprocess.PIPE) time.sleep(1) else: if total_attempt > 5: cmdwrite.status('local client maxed out attempt') return else: break while True: early_exit = False if conn.poll(20) == True: response = json.loads(conn.recv()) if 'local_debug_message' in response.keys(): display_debug_message(response) elif 'status' in response.keys(): early_exit = display_judge_status(response, problem_id) else: cmdwrite.client_debug('server time out') if early_exit == True: break conn.close()
class SocketDelegate(Delegate): def __init__(self): self._client = None self._queue = [] self._poll_interval = 0.001 self._update_interval = 0.001 self._time_since_last_update = 0 def connect(self, port: int): self._client = Client(('localhost', port)) Thread(target=self._run, daemon=True).start() def on_custom_event(self, code: str, event: str, data: Any) -> Any: pass def _run(self): while True: time.sleep(self._poll_interval) self._time_since_last_update += self._poll_interval self._read_events() if self._time_since_last_update >= self._update_interval: self._time_since_last_update = 0 self._update() def _read_events(self): while self._client.poll(): self._queue.append(self._client.recv()) def _update(self): for t, event, data in self._queue: if t == 'c': self._client.send(self.check_event(event, data)) elif t == 'e': self.on_event(event, data) else: result = self.on_custom_event(t, event, data) if result is not None: self._client.send(result) self._queue.clear()
def main(): pipe = Client(sys.argv[1]) proto, target = pipe.recv() try: s = lightblue.socket(proto) s.connect(target) s.settimeout(1 / 100.) pipe.send(None) while True: if pipe.poll(): b = pipe.recv() print "sending", b s.send(b) try: o = s.recv(4096) pipe.send(o) except Exception, err: pass except Exception, err: pipe.send(err) raise
class Client(object): """Send a command to a listener process. Args: address (tuple/object): The address to be used by the bound socket or named pipe (EX: "IP Address", Port) family (socket.family/str)[None]: Type of socket or named pipe to use (EX: 'AF_INET', 'AF_UNIX', 'AF_PIPE') authkey (bytes)[None]: The secret key (password) for an HMAC-based authentication challenge. No auth if None Raises: AuthenticationError: If authkey is given and authentication fails """ MY_IP = MY_IP get_local_addrs = staticmethod(get_local_addrs) def __init__(self, address, family=None, authkey=None): """Create the client connection to another process. Args: address (tuple/object): The address to be used by the bound socket or named pipe (EX: "IP Address", Port) family (socket.family/str)[None]: Type of socket or named pipe to use (EX: 'AF_INET', 'AF_UNIX', 'AF_PIPE') authkey (bytes)[None]: The secret key (password) for an HMAC-based authentication challenge. No auth if None Raises: AuthenticationError: If authkey is given and authentication fails """ if isinstance(authkey, str): authkey = authkey.encode('utf-8') # Connection settings self.address = address self.family = family # self.authkey = authkey # Probably not good to save the authkey self.client = None self.connect(address, family, authkey) def connect(self, address=None, family=None, authkey=None): """Connect the multiprocessing client.""" if address is not None: self.address = address if family is not None: self.family = family self.client = MpClient(self.address, family=self.family, authkey=authkey) # Why is this a function!? F def msg_handler(self, cmd): """Handle the response after sending a command. This is where you would receive an ack or nack from the socket. Args: cmd (object): Command that was sent. Returns: success (bool): True if the message sent successfully! """ return True @staticmethod def recv_socket(sock): """Receive data from the connection.""" return sock.recv() @staticmethod def send_socket(sock, data): """Send data to the connection.""" sock.send(data) def send(self, cmd): """Send the data to the server.""" self.send_socket(self.client, cmd) return self.msg_handler(cmd) def recv(self): """Receive data from the server.""" return self.recv_socket(self.client) def poll(self, timeout=0.0): """Whether there is any input available to be read""" return self.client.poll(timeout=timeout) @property def closed(self): """True if the connection is closed""" return self.client.closed @property def readable(self): """True if the connection is readable""" return self.client.readable @property def writable(self): """True if the connection is writable""" return self.client.writable def fileno(self): """File descriptor or handle of the connection""" return self.client.fileno() def close(self): """Close the connection""" time.sleep(0.01) # Could be sending a message self.client.close() stop = close def __enter__(self): return self def __exit__(self, exc_type, exc_value, exc_tb): self.close() return False
class _NativeAppClient: def __init__(self): self.__rawPluginOutputStream = Subject() self.__connectionStatusStream = Subject() self.__speechStream = with_latest_from( self.__rawPluginOutputStream.pipe( map(lambda dehydratedMsgDict: rehydrateMessage( dehydratedMsgDict))), self.__connectionStatusStream, ).pipe( map(lambda combinedTuple: { **combinedTuple[0], **combinedTuple[1] }), merge(self.__connectionStatusStream), ) # TODO - need to trigger NVDA to startup, if it isn't already # - first need to check if NVDA is installed + if the plugin is asyncio.create_task(self.__startListeningForOutput()) @property def speechStream(self): return self.__speechStream # TODO - add streams for other vitals checks. # - Consumer to choose frequencies? # - Is NVDA running? async def __startListeningForOutput(self): await self.__connectToOutputServer() asyncio.create_task(self.__pollForOutput()) async def __connectToOutputServer(self): self.__serverConnection = None # Remove a potential previous, broken connection await asyncio.sleep(_NVDA_EXPECTED_STARTUP_TIME + _PADDING_FOR_NVDA_STARTUP_TIME) while True: try: self.__serverConnection = Client( _address, _authkey) # Seems to wait ~3s before exception is raised self.__connectionStatusStream.on_next( {OutputKeys.IS_CONNECTED: True}) break except ConnectionRefusedError: # If the listener hasn't been setup yet self.__connectionStatusStream.on_next( {OutputKeys.IS_CONNECTED: False}) await asyncio.sleep(0) async def __pollForOutput(self): while True: if self.__serverConnection and self.__serverConnection.poll(): try: dehydratedMsgDict = self.__serverConnection.recv() self.__rawPluginOutputStream.on_next(dehydratedMsgDict) except ConnectionResetError: # If NVDA/the plugin is terminated self.__connectionStatusStream.on_next( {OutputKeys.IS_CONNECTED: False}) await self.__connectToOutputServer() # In case it restarts else: await asyncio.sleep(0)
def train(self): while True: while not self.received: try: send_address = ('100.97.66.131', 8000) conn = Client(send_address, authkey=b'nasnet') if conn.poll(2): [self.network] = conn.recv() self.total_recv += 1 conn.close() self.received = True self.dump_client() print("RECEIEVE:=>", self.network) print("RECEIEVE:=>", " total_send:", self.total_send, " total_recv:", self.total_recv) self.print_client_status() except Exception as e: print(e) print(traceback.format_exc()) print("client recv error") if self.received: print("prepare training the network:", self.network) network = np.array(self.network, dtype='int') network = network.tolist() net = gen_code_from_list(network, node_num=7) #TODO: change it to 7 net_str = json.dumps(network) if net_str in self.accuracy_trace: self.acc = self.accuracy_trace[net_str] else: genotype_net = translator( [net, net], max_node=7) #TODO: change it to 7 print("--" * 15) print(genotype_net) print("training the above network") print("--" * 15) self.acc = train_client.run(genotype_net, epochs=600, batch_size=200) self.accuracy_trace[net_str] = self.acc self.dump_acc_trace() #TODO: train the actual network #time.sleep(random.randint(2, 5) ) while self.received: try: recv_address = ('100.97.66.131', 8000) conn = Client(recv_address, authkey=b'nasnet') network_str = json.dumps(np.array(network).tolist()) conn.send([self.client_name, network_str, self.acc]) self.total_send += 1 print("SEND:=>", self.network, self.acc) self.network = [] self.acc = 0 self.received = False self.dump_client() print("SEND:=>", " total_send:", self.total_send, " total_recv:", self.total_recv) conn.close() except Exception as e: print(e) print(traceback.format_exc()) print("client send error, reconnecting")
# take the highest priority process QueueTop = priority_queue.pop() TopRequestID = -QueueTop[1] # send data to processing module conn_pq2p.send(RequestID_to_json_data[TopRequestID]) # remove top request from map RequestID_to_json_data.pop(TopRequestID) # after sending data, processor will be busy processor_busy = True # if priority_queue is empty and processor no busy and connection from adapter has already been terminated # then close the communication to processing module elif not conn_a2pq_active: conn_pq2p.send("terminate") conn_pq2p.close() running = False break if conn_pq2p.poll(): msg = conn_pq2p.recv() # make it non blocking if (msg == "free"): processor_busy = False # priority_queue will keep on receiving requests from adapter.py # but will send data to processor only if it is free to do so # since it is a client server relation # hence feedbacks can be used from processing module # to check whether it is available to process next query # if yes then send next query
class PlcMonitorThread(Thread): def __init__(self, dev, plc): Thread.__init__(self) self.dev = dev self.plc = plc self.running = False self.listener_ready = False def __create_connection(self): tmp_base = utils.get_tmp_base_path_from_mkfile() plc_base_path = os.path.join(tmp_base, self.plc.name) address = os.path.join(plc_base_path, 'plc_socket') # Wait until listener is ready # TODO: Refactoring needed (anti-pattern) while not self.listener_ready: try: self.conn = Client(address) self.listener_ready = True except socket.error as e: # Check if listener is not yet ready if e.errno == errno.ENOENT: # socket.error: [Errno 2] No such file or directory sleep(1) else: logger.exception("Unknown socket error occurred.") def __monitor(self): self.__create_connection() self.conn.send(MonitorMessage(self.dev.plc_vars_map.keys())) def __close(self): self.conn.close() def run(self): self.__monitor() while not self.dev.shutdown_event.is_set(): try: if self.conn.poll(): result = self.conn.recv() if isinstance(result, UnknownPlcTagException): raise UnknownPlcTagException(result) elif isinstance(result, TerminateMessage): logger.info("Terminating monitoring of '{}'.".format( self.dev.name)) break elif isinstance(result, MonitorResponseMessage): dev_var_idx = self.dev.plc_vars_map[result.name] var = self.dev.vars[dev_var_idx]['value'] if type(var) is int: self.dev.vars[dev_var_idx]['value'] = int( result.value) elif type(var) is bool: self.dev.vars[dev_var_idx]['value'] = result.value else: raise RuntimeError( 'Unsupported type \'{}\'.'.format(type(var))) else: logger.error("Received unexpected message type '%s'.", type(result)) except EOFError: logger.exception("Received EOF.") break except UnknownPlcTagException: logger.exception("Unknown PLC Tag.") self.__close()
class Rexec(BaseObj): """Rexec object Rexec() -> New remote procedure object Arguments: servername: Name or IP address of remote server logfile: Name of logfile to create on remote server Usage: from nfstest.rexec import Rexec # Function to be defined at remote host def add_one(n): return n + 1 # Function to be defined at remote host def get_time(delay=0): time.sleep(delay) return time.time() # Create remote procedure object x = Rexec("192.168.0.85") # Define function at remote host x.rcode(add_one) # Evaluate the expression calling add_one() out = x.reval("add_one(67)") # Run the function with the given argument out = x.run("add_one", 7) # Run built-in functions import time out = x.run(time.time) # Import libraries and symbols x.rimport("time", ["sleep"]) x.run("sleep", 2) # Define function at remote host -- since function uses the # time module, this module must be first imported x.rimport("time") x.rcode(get_time) # Evaluate the expression calling get_time() out = x.reval("get_time()") # Run the function with the given argument out = x.run("get_time", 10) # Open file on remote host fd = x.run(os.open, "/tmp/testfile", os.O_WRONLY|os.O_CREAT|os.O_TRUNC) count = x.run(os.write, fd, "hello there\n") x.run(os.close, fd) # Use of positional arguments out = x.run('get_time', 2) # Use of named arguments out = x.run('get_time', delay=2) # Use of NOWAIT option for long running functions so other things # can be done while waiting x.run('get_time', 2, NOWAIT=True) while True: # Poll every 0.1 secs to see if function has finished if x.poll(0.1): # Get results out = x.results() break """ def __init__(self, servername="", logfile=None): """Constructor Initialize object's private data. servername: Host name or IP address of host where remote server will run logfile: Pathname of log file to be created on remote host [Default: "/dev/null"] """ global PORT self.pid = None self.conn = None self.process = None if logfile is None: # Default log file logfile = "/dev/null" if servername in ["", "localhost", "127.0.0.1"]: # Start bare-bones server locally self.remote = False self.pid = os.fork() if self.pid == 0: # This is the child process bare_server(PORT, logfile) os._exit(0) else: # Start bare-bones server on remote host self.remote = True server_code = "".join(inspect.getsourcelines(bare_server)[0]) server_code += 'bare_server(%d, "%s")' % (PORT, logfile) server_code = re.sub(r"'", r"\\'", server_code) cmdlist = ["ssh", "-t", servername, "python -c '%s'" % server_code] self.process = Popen(cmdlist, shell=False, stdin=PIPE, stdout=PIPE, stderr=PIPE) # Connect to remote server address = (servername, PORT) self.conn = Client(address) PORT += 1 # Execute main server on remote host proc_code = "".join(inspect.getsourcelines(proc_requests)[0]) proc_code += "proc_requests(fd, conn)" self.conn.send(proc_code) def __del__(self): """Destructor""" if self.conn: # Send command to exit main loop self.conn.send('close') self.conn.close() # Wait for remote server to finish if self.pid: os.waitpid(self.pid, 0) elif self.process: self.process.wait() def _send_cmd(self, cmd, *kwts, **kwds): """Internal method to send commands to remote server""" nowait = kwds.pop('NOWAIT', False) self.conn.send({'cmd': cmd, 'kwts': kwts, 'kwds': kwds}) if nowait: # NOWAIT option is specified, so return immediately # Use poll() method to check if any data is available # Use results() method to get pending results from function return return self.results() def wait(self, objlist=None, timeout=0): """Return a list of Rexec objects where data is available to be read objlist: List of Rexec objects to poll, if not given use current object timeout: Maximum time in seconds to block, if timeout is None then an infinite timeout is used """ ret = [] if objlist is None: # Use current object as default objlist = [self] for obj in objlist: if obj.poll(timeout): ret.append(obj) # Just check all other objects if they are ready now timeout = 0 return ret if len(ret) else None def poll(self, timeout=0): """Return whether there is any data available to be read timeout: Maximum time in seconds to block, if timeout is None then an infinite timeout is used """ return self.conn.poll(timeout) def results(self): """Return pending results""" out = self.conn.recv() if isinstance(out, Exception): raise out return out def rexec(self, expr): """Execute statement on remote server""" return self._send_cmd('exec', expr) def reval(self, expr): """Evaluate expression on remote server""" return self._send_cmd('eval', expr) def run(self, *kwts, **kwds): """Run function on remote server The first positional argument is the function to be executed. All other positional arguments and any named arguments are treated as arguments to the function """ return self._send_cmd('run', *kwts, **kwds) def rcode(self, code): """Define function on remote server""" codesrc = "".join(inspect.getsourcelines(code)[0]) self.rexec(codesrc) def rimport(self, module, symbols=[]): """Import module on remote server module: Module to import in the remote server symbols: If given, import only these symbols from the module """ # Import module if len(symbols) == 0: self.rexec("import %s" % module) symbols = [module] else: self.rexec("from %s import %s" % (module, ",".join(symbols))) # Make all symbols global for item in symbols: self.rexec("globals()['%s']=locals()['%s']" % (item, item))
class HumanLabelWrapper(gym.Wrapper): def __init__(self, env_id, env, online_blocking_mode="action_replacement", reward_scale=1.0, catastrophe_reward=0.0, **_): super().__init__(env) self.env_id = env_id self.conn = None self.current_frame = None self.episode_info_common = {} self.online_blocking_mode = online_blocking_mode self.catastrophe_reward = catastrophe_reward self.reward_scale = reward_scale self.episode_block_count = 0 if isinstance(self.env.unwrapped, gym.envs.atari.atari_env.AtariEnv): self.episode_info_common["action_set"] = self.env.unwrapped._action_set # todo(girish): assumes human_feedback.py is running and listening on port 6666 logger.info('Opening connection...') # open connection on 6666 (won't support multiple workers) address = ('localhost', 6666) self.conn = Client(address, authkey=b'no-catastrophes-allowed') time.sleep(2.0) def _reset(self): self.__send_init_msg() observation = self.env.reset() # note - when run with FrameSaverWrapper, render happens twice (not sure if cached) image = self.env.render("rgb_array") self.current_frame = Frame(observation, None, 0, image, {}) return observation def __send_init_msg(self): logger.info('Sending initial message...') init_msg_dict = { 'msg': 'init' } # todo - don't get episode num in human_feedback.py from this dict init_msg_dict.update(self.episode_info_common) init_msg_dict['env_id'] = self.env_id self.conn.send(init_msg_dict) def __send_frame_to_human(self, frame, proposed_action): """ Sends a frame with a proposed action to the human for feedback. """ # if proposed_action is not None and frame.action is not None: frame.action = proposed_action frame.info["frame/action/proposed_action"] = proposed_action # info["frame/action/real_action"] = real_action proposed_action_meaning = ACTION_MEANING[self.episode_info_common["action_set"][ proposed_action]] logger.info('Sending frame. Proposed action: {}'.format(proposed_action_meaning)) self.conn.send(frame) def __handle_feedback(self, msg, proposed_action): """Given the proposed action and the human feedback message Returns: - a tuple (action, label). """ logger.info('online step, received: {}'.format(msg)) real_action = proposed_action label = None # parse feedback message if isinstance(msg, dict): label = 'b' if self.online_blocking_mode == "action_pruning": return None, label real_action = msg['action'] proposed_action_meaning = ACTION_MEANING[self.episode_info_common["action_set"][ proposed_action]] real_action_meaning = ACTION_MEANING[self.episode_info_common["action_set"][ real_action]] logger.info('proposed action {} was blocked, replacing with action {}' .format(proposed_action_meaning, real_action_meaning)) return real_action, label def _step(self, proposed_action): self.__send_frame_to_human(self.current_frame, proposed_action=proposed_action) if not self.conn.poll(100.0): raise EnvironmentError("Failed to receive message!") msg = self.conn.recv() real_action, label = self.__handle_feedback(msg, proposed_action) if real_action is None: logger.info('Pruning action') observation, reward, done, info = self.current_frame.observation, \ self.catastrophe_reward / self.reward_scale, \ False, \ self.current_frame.info info["frame/reward"] = reward info["frame/action/proposed_action"] = proposed_action info["frame/action/real_action"] = real_action else: logger.info( 'Taking action: {}' .format(ACTION_MEANING[self.episode_info_common["action_set"][real_action]])) observation, reward, done, info = self.env.step(real_action) image = self.env.render("rgb_array") if real_action is not proposed_action: reward += self.catastrophe_reward / self.reward_scale logger.info('Action was changed.') self.episode_block_count += 1 logger.info('Updated episode block count to {}'.format(self.episode_block_count)) # self.__send_frame_to_human(self.current_frame, proposed_action=real_action) # self.conn.recv() info["frame/reward"] = reward # initialize current frame for next step self.current_frame = Frame(observation, None, reward, image, info.copy()) info["frame/action/proposed_action"] = proposed_action info["frame/action/real_action"] = real_action if done: logger.info('Sending done message') self.conn.send({'msg': 'done'}) info["global/episode_block_count"] = self.episode_block_count self.episode_block_count = 0 info["frame/action/label"] = label return observation, reward, done, info def _close(self): if self.conn is not None: print('Closing connection on {host}:{port}'.format(host='localhost', port=6666)) self.conn.send({'msg': 'close'}) self.conn.close()
class Nest(Process): """ Represents a group of ants. It is realised by a Process class. A host run one nest per CPU core by default. Use the --nextcount argument to change it. """ def __init__(self, address: str, port: int, name='noname'): super(Nest, self).__init__(name=name) try: # open connection towards the colony self._conn = Client(address=(address, port), family='AF_INET') self._stopevent = Event() self._startevent = Event() self._ants = [] # load simulation Queen # with open("examples/simple.py", "r") as f: # exec(f.read()) # scheduler used for scheduling hatch events of Eggs. self._scheduler = sched.scheduler(time.time, time.sleep) self.start() except ConnectionRefusedError as err: print("Cannot connect to Colony on %s:%d: '%s', terminating." % (address, port, err)) sys.exit(1) except Exception as e: self._log(e) def addant(self, ant: Ant): """ Keep track of hatched Ants. """ # assert issubclass(ant, type(Ant)), "Only ants can be added to nest, I got '%s'" % type(ant) #TODO: fix this # set Ant's conn for logging ant.conn = self._conn self._ants.append(ant) ant.start() def run(self): """ Manage egg hatching, communication with Colony. """ self._log("started, max open files: '%d %d'" % resource.getrlimit(resource.RLIMIT_NOFILE)) try: # keep checking commands till stopevent while not self._stopevent.isSet(): # poll till the next event, or for 0.1 seconds, if not yet executing. polltime = self._scheduler.run( blocking=False) if self._startevent.isSet() else 0.1 # remember the start of pooling. If waiting during polling is interrupted, the remaining time must be # polled again. pollstart = time.time() # terminate, if there are no more eggs to hatch (Scheduler.run() returns None) if polltime is None: self._stopevent.set() # do polling, repeate with the remaining wait time, if interrupted by a message. while not self._stopevent.isSet() and self._conn.poll( timeout=max(0, polltime)): o = self._conn.recv() # if isinstance(o, Egg): #TODO figure out, why this is not working if o.__class__.__name__ == 'Egg': self._scheduler.enter(delay=o.delay, priority=1, action=o.hatch, argument=( self, self._conn, )) self._log("egg with larv '%s' to hatch at %.2f" % (o.larv, o.delay)) elif isinstance(o, Cmd): if o.isexecute(): self._startevent.set() elif o.isterminate(): self._terminate() # calculate remaining wait time polltime -= time.time() - pollstart # wait for still running ants to finish for ant in self._ants: if ant.is_alive(): ant.join() self._ants.remove(ant) except KeyboardInterrupt: self._log("interrupted.") self._conn.send(Cmd.terminated()) self._conn.close() sys.exit(0) def _terminate(self): self._stopevent.set() # termiante all running ants for ant in self._ants: if ant.is_alive(): ant.terminate() def _log(self, logstring): """ Use remote logging on Colony. """ self._conn.send( Msg("%s '%s': %s" % (self.__class__.__name__, self.name, logstring)))
class Rexec(BaseObj): """Rexec object Rexec() -> New remote procedure object Arguments: servername: Name or IP address of remote server logfile: Name of logfile to create on remote server sudo: Run remote server as root Usage: from nfstest.rexec import Rexec # Function to be defined at remote host def add_one(n): return n + 1 # Function to be defined at remote host def get_time(delay=0): time.sleep(delay) return time.time() # Create remote procedure object x = Rexec("192.168.0.85") # Define function at remote host x.rcode(add_one) # Evaluate the expression calling add_one() out = x.reval("add_one(67)") # Run the function with the given argument out = x.run("add_one", 7) # Run built-in functions import time out = x.run(time.time) # Import libraries and symbols x.rimport("time", ["sleep"]) x.run("sleep", 2) # Define function at remote host -- since function uses the # time module, this module must be first imported x.rimport("time") x.rcode(get_time) # Evaluate the expression calling get_time() out = x.reval("get_time()") # Run the function with the given argument out = x.run("get_time", 10) # Open file on remote host fd = x.run(os.open, "/tmp/testfile", os.O_WRONLY|os.O_CREAT|os.O_TRUNC) count = x.run(os.write, fd, "hello there\n") x.run(os.close, fd) # Use of positional arguments out = x.run("get_time", 2) # Use of named arguments out = x.run("get_time", delay=2) # Use of NOWAIT option for long running functions so other things # can be done while waiting x.run("get_time", 2, NOWAIT=True) while True: # Poll every 0.1 secs to see if function has finished if x.poll(0.1): # Get results out = x.results() break # Create remote procedure object as a different user # First, run the remote server as root x = Rexec("192.168.0.85", sudo=True) # Then set the effective user id x.run(os.seteuid, 1000) """ def __init__(self, servername=None, logfile=None, sudo=False, sync_timeout=0.1): """Constructor Initialize object's private data. servername: Host name or IP address of host where remote server will run [Default: None (run locally)] logfile: Pathname of log file to be created on remote host [Default: None] sudo: Run remote procedure server as root [Default: False] sync_timeout: Timeout used for synchronizing the connection stream [Default: 0.1] """ global PORT self.pid = None self.conn = None self.process = None self.remote = False self.servername = servername self.logfile = logfile self.sudo = sudo self.sync_timeout = sync_timeout if os.getuid() == 0: # Already running as root self.sudo = True sudo = False if not sudo and servername in [None, "", "localhost", "127.0.0.1"]: # Start remote server locally via fork when sudo is not set servername = "" self.pid = os.fork() if self.pid == 0: # This is the child process RemoteServer(PORT, self.logfile).start() os._exit(0) else: # Start server on remote host or locally if sudo is set server_code = IMPORTS server_code += "".join(inspect.getsourcelines(RemoteServer)[0]) server_code += "RemoteServer(%d, %r).start()\n" % (PORT, self.logfile) # Execute minimal python script to execute the source code # given in standard input pysrc = "import sys; exec(sys.stdin.read(%d))" % len(server_code) cmdlist = ["python", "-c", repr(pysrc)] if sudo: cmdlist.insert(0, "sudo") if servername not in [None, "", "localhost", "127.0.0.1"]: # Run remote process via ssh cmdlist = ["ssh", servername] + cmdlist self.process = Popen(cmdlist, shell=False, stdin=PIPE) self.remote = True else: # Run local process via the shell servername = "" self.process = Popen(" ".join(cmdlist), shell=True, stdin=PIPE) # Send the server code to be executed via standard input self.process.stdin.write(server_code) # Connect to remote server address = (servername, PORT) self.conn = Client(address) PORT += 1 def __del__(self): """Destructor""" self.close() def close(self): """Close connection to remote server""" if self.conn: # Send command to exit main loop self.conn.send("close") self.conn.close() self.conn = None # Wait for remote server to finish if self.pid: os.waitpid(self.pid, 0) self.pid = None elif self.process: self.process.wait() self.process = None def _send_cmd(self, cmd, *kwts, **kwds): """Internal method to send commands to remote server""" nowait = kwds.pop("NOWAIT", False) self.conn.send({"cmd": cmd, "kwts": kwts, "kwds": kwds}) if nowait: # NOWAIT option is specified, so return immediately # Use poll() method to check if any data is available # Use results() method to get pending results from function return return self.results() def wait(self, objlist=None, timeout=0): """Return a list of Rexec objects where data is available to be read objlist: List of Rexec objects to poll, if not given use current object timeout: Maximum time in seconds to block, if timeout is None then an infinite timeout is used """ ret = [] if objlist is None: # Use current object as default objlist = [self] for obj in objlist: if obj.poll(timeout): ret.append(obj) # Just check all other objects if they are ready now timeout = 0 return ret if len(ret) else None def poll(self, timeout=0): """Return whether there is any data available to be read timeout: Maximum time in seconds to block, if timeout is None then an infinite timeout is used """ return self.conn.poll(timeout) def results(self): """Return pending results""" while True: out = self.conn.recv() if isinstance(out, Exception): raise out elif out is None and self.poll(self.sync_timeout): # Try to re-sync when recv() returns None and there is # still data in the buffer continue return out def rexec(self, expr): """Execute statement on remote server""" return self._send_cmd("exec", expr) def reval(self, expr): """Evaluate expression on remote server""" return self._send_cmd("eval", expr) def run(self, *kwts, **kwds): """Run function on remote server The first positional argument is the function to be executed. All other positional arguments and any named arguments are treated as arguments to the function """ return self._send_cmd("run", *kwts, **kwds) def rcode(self, code): """Define function on remote server""" codesrc = "".join(inspect.getsourcelines(code)[0]) self.rexec(codesrc) def rimport(self, module, symbols=[]): """Import module on remote server module: Module to import in the remote server symbols: If given, import only these symbols from the module """ # Import module if len(symbols) == 0: self.rexec("import %s" % module) symbols = [module] else: self.rexec("from %s import %s" % (module, ",".join(symbols))) # Make all symbols global for item in symbols: self.rexec("globals()['%s']=locals()['%s']" % (item, item))
class ClientUserInterface: """ This Class is the Interface for the user script to the service script. """ def __init__(self, locked): """ This initializes the Interface between the service script and the user script. Args: locked: the callback function to check if the user screen is locked """ try: self.__connection = Client(ADDRESS) except ConnectionRefusedError: notify_window = Tk() notify_window.title("Smock") label_username = Label( notify_window, text= "Service script isn't running.\n please start the service script" "\n and restart the programm then.") btn_ok = Button(notify_window, text="OK", command=notify_window.destroy) label_username.pack() btn_ok.pack() notify_window.mainloop() self.__mutex = Lock() self.gui = None self.is_locked = locked def polling(self): """ Checks if a message is sent by the service script. If a message is in the buffer, it reads the command and reacts according to the given command. """ while True: time.sleep(0.1) refresh = False self.__mutex.acquire(True) if self.__connection.poll(): data = self.__connection.recv() if data == str(Command.COMPUTER_STATUS): multi_user = self.__connection.recv() ret = str(self.is_locked(multi_user)) self.__connection.send(ret) elif data == Codes.GUI_UPDATE: uid = self.__connection.recv() self.update_gui(uid) elif data == Codes.FIND_SERIAL_DEVICES: if self.__connection.recv() == "True": Gui.notify("Das Gerät ist bereits angeschlossen") else: Gui.notify( "Es wurde kein Smock Gerät gefunden.\n" "Vergewissern Sie sich, dass das Gerät angeschlossen ist.\n" ) elif data == Codes.REFRESH: refresh = True self.__mutex.release() if refresh: self.gui.refresh_list() def check_if_user_exists(self, name): """ Asks the service script if the user exists Args: name: the name of the user that is asked for Returns: True, when user exists, otherwise False """ self.__mutex.acquire(True) self.__connection.send(Codes.USER_EXISTING) self.__connection.send(name) data = self.__connection.recv() self.__mutex.release() return data def get_username(self, index): """ Asks the service script for the username of a user and returns the username. Args: index: index of the user in the list Returns: username of the user """ self.__mutex.acquire(True) self.__connection.send(Codes.GET_USERNAME) self.__connection.send(str(index)) data = self.__connection.recv() self.__mutex.release() return data def get_uid(self, index): """ Asks the service script for the uid of a user and returns the uid. Args: index: index of the user in the list Returns: uid of the user """ self.__mutex.acquire(True) self.__connection.send(Codes.GET_UID) self.__connection.send(str(index)) data = self.__connection.recv() self.__mutex.release() return data def set_new_data(self, index, user_name, user_pwd): """ Sends the new data of a user to the service script. This will edit the user data in the service script. Args: index: index of user in list user_name: username of user user_pwd: password of user """ self.__mutex.acquire(True) self.__connection.send(Codes.SET_NEW_DATA) self.__connection.send(str(index) + ";" + user_name + ";" + user_pwd) # wait for acknowledge self.__connection.recv() self.__mutex.release() def write_to_controller(self, msg): """ Sends the service script the code to forward a message to the serial device and the message that is send to the serial device Args: msg: message that is send to the service script """ self.__mutex.acquire(True) self.__connection.send(Codes.WRITE_TO_CONTROLLER) self.__connection.send(msg) self.__mutex.release() def add_user(self, username, pwd, uid): """ Sends the service script the information of adding a user and the data of the user. Args: username: username of the new user pwd: password of the username uid: uid of the username """ self.__mutex.acquire(True) self.__connection.send(Codes.ADD_USER) self.__connection.send(username + ";" + pwd + ";" + uid) self.__mutex.release() def delete_user(self, index): """ Sends the service script the information to delete a user and the index of that user. Args: index: the index of the user in the list """ self.__mutex.acquire(True) self.__connection.send(Codes.DELETE_USER) self.__connection.send(str(index)) self.__mutex.release() def get_number_of_users(self): """ Gets the number of users that are in the user list from the service script Returns: The number of users that exist """ self.__mutex.acquire(True) self.__connection.send(Codes.GET_NUMBER_OF_USERS) data = int(self.__connection.recv()) self.__mutex.release() return data def find_serial_device(self): """ Sends code of finding a serial device to the service script """ self.__mutex.acquire(True) self.__connection.send(Codes.FIND_SERIAL_DEVICES) self.__mutex.release() def update_gui(self, uid): """ updates the uid label for the gui, when the user is in the add user window. Args: uid: the uid that replaces the old label """ self.gui.update_uid_label(uid)
def main(args): # Connect to aruco server for pose estimates try: conn = Client(('localhost', 6000), authkey=b'secret password') except ConnectionRefusedError: print('Could not connect to aruco server for pose estimates') return # Create action executor for the physical robot action_executor = vector_action_executor.VectorActionExecutor( args.robot_index) # Create env config_path = args.config_path if config_path is None: config_path = utils.select_run() if config_path is None: print('Please provide a config path') return cfg = utils.read_config(config_path) kwargs = {'num_cubes': args.num_cubes} if args.debug: kwargs['use_gui'] = True cube_indices = list(range(args.num_cubes)) env = utils.get_env_from_cfg(cfg, physical_env=True, robot_index=action_executor.robot_index, cube_indices=cube_indices, **kwargs) env.reset() # Create policy policy = utils.get_policy_from_cfg(cfg, env.get_action_space()) # Debug visualization if args.debug: cv2.namedWindow('out', cv2.WINDOW_NORMAL) #cv2.resizeWindow('out', 960, 480) try: while True: # Get new pose estimates poses = None while conn.poll(): # ensure up-to-date data poses = conn.recv() if poses is None: continue # Update poses in the simulation env.update_poses(poses) # Get new action state = env.get_state() if action_executor.is_action_completed() and args.debug: action, info = policy.step(state, debug=True) # Visualize assert not cfg.use_steering_commands output = info['output'].cpu().numpy() cv2.imshow( 'out', utils.get_state_and_output_visualization( state, output)[:, :, ::-1]) cv2.waitKey(1) else: action, _ = policy.step(state) # Run selected action through simulation try_action_result = env.try_action(action) if action_executor.is_action_completed(): # Update action executor action_executor.update_try_action_result(try_action_result) # Run action executor action_executor.step(poses) finally: action_executor.disconnect()
class RemoteAgent: def __init__(self, with_adaptation=True, connection_retries=100): atexit.register(self.terminate) self._log = logging.getLogger(self.__class__.__name__) sock_file = tempfile.mktemp() cmd = [ sys.executable, # path to the current python binary str((pathlib.Path(__file__).parent.parent / "zoo" / "run_agent.py").absolute().resolve()), sock_file, ] if with_adaptation: cmd.append("--with_adaptation") self._log.debug(f"Spawning remote agent proc: {cmd}") self._agent_proc = subprocess.Popen(cmd) self._conn = None for i in range(connection_retries): # Waiting on agent to open it's socket. try: self._conn = Client(sock_file, family="AF_UNIX") break except FileNotFoundError: self._log.debug( f"RemoteAgent retrying connection to agent in: attempt {i}" ) time.sleep(0.1) if self._conn is None: raise RemoteAgentException("Failed to connect to remote agent") def __del__(self): self.terminate() def send_observation(self, obs): self._conn.send({"type": "obs", "payload": obs}) def recv_action(self, timeout=None): if self._conn.poll(timeout): try: return self._conn.recv() except ConnectionResetError as e: self.terminate() raise e else: return None def start(self, agent_spec: AgentSpec): # send the AgentSpec to the agent runner self._conn.send( # We use cloudpickle only for the agent_spec to allow for serialization of lambdas { "type": "agent_spec", "payload": cloudpickle.dumps(agent_spec) }) def terminate(self): if self._agent_proc: if self._conn: self._conn.close() if self._agent_proc.poll() is not None: self._agent_proc.kill() self._agent_proc.wait() self._agent_proc = None
class BinDbgSession: def __init__(self, bv): self.bv = bv self.conn = None self.windbg_proc = None self.regs = None self.ip = None self.bps = set() if not 'proc_args' in self.bv.session_data: self.bv.session_data['proc_args'] = '' if not 'pipe' in self.bv.session_data: self.bv.session_data['pipe'] = get_text_line_input( "Enter BinDbg pipe name:", "Start BinDbg session", ) t = threading.Thread(target=self.event_loop) t.setDaemon(True) t.start() def stop(self, reason): if 'bindbg' in self.bv.session_data: if self.bv.session_data['bindbg'].windbg_proc: # close windbg pid = self.bv.session_data['bindbg'].windbg_proc.pid for hwnd in get_hwnds_for_pid(pid): win32gui.PostMessage(hwnd, win32con.WM_CLOSE, 0, 0) del self.bv.session_data['bindbg'] if self.conn: self.conn.close() for bp in self.bps.copy(): self.delete_bp(bp) if self.ip: self.highlight(self.ip, no_color) self.ip = None print(reason) def send(self, command, **params): try: self.conn.send((command, params)) except IOError: return self.stop("Lost connection to WinDbg") def connect(self): try: self.conn = Client('\\\\.\\pipe\\' + self.bv.session_data['pipe']) return True except WindowsError: return False def start_windbg(self): bindbg_dir = os.path.dirname(os.path.abspath(__file__)) windbg_ext = os.path.join(bindbg_dir, 'windbg.py') exe_path = self.bv.file.filename.replace('.bndb', '.exe') exe_dir = os.path.dirname(exe_path) # determine whether we should run x86 or x64 windbg if self.bv.arch.name == 'x86': windbg_path = os.path.join(dbg_dir, 'x86', 'windbg.exe') else: windbg_path = os.path.join(dbg_dir, 'x64', 'windbg.exe') # save command as a script due to nested double quotes windbg_cmd = '.load "{}"; !py -2 --global "{}" "{}"'.format( pykd_path.replace('\\', '\\\\'), windbg_ext.replace('\\', '\\\\'), self.bv.session_data['pipe']) script_path = os.path.join(bindbg_dir, 'cmd.txt') with open(script_path, "w") as script: script.write(windbg_cmd) # open WinDbg without blocking self.windbg_proc = subprocess.Popen( [ windbg_path, '-W', exe_path, # create a named workspace for this exe '-QY', # automatically save workspace to preserve BPs & layout '-c', '$$<' + script_path, # run windbg_cmd exe_path, self.bv.session_data['proc_args'] ], cwd=exe_dir, ) def event_loop(self): # attempt to connect to already open windbg session if not self.connect(): # if that fails, start windbg and attempt to connect again self.start_windbg() while not self.connect(): # wait while windbg loads time.sleep(2) print("Connected to WinDbg") self.send('sync') while True: try: if self.conn.poll(poll_time): self.process(self.conn.recv()) except IOError as e: if e.errno == 6: # stop() has already been called return return self.stop("Lost connection to WinDbg") def process(self, data): print(data) cmd, params = data if 'bp' in cmd: addr = params['addr'] + self.bv.start if self.bv.is_valid_offset(addr): if cmd == 'set_bp': self.set_bp(addr) elif cmd == 'delete_bp': self.delete_bp(addr) elif cmd == 'bp_hit': self.bp_hit(addr, params['regs']) elif cmd == 'set_ip': ip = params['ip'] + self.bv.start if self.bv.is_valid_offset(ip): self.set_ip(ip, params['regs']) elif cmd == 'vtable': ip = params['ip'] + self.bv.start target = params['target'] + self.bv.start object = params[ 'object'] + self.bv.start if 'object' in params else None instr = params['instr'] self.vtable(ip, target, object, instr) def highlight(self, addr, color): func = self.bv.get_functions_containing(addr)[0] func.set_auto_instr_highlight(addr, color) def set_ip(self, addr, regs=None, came_from_binja=False): if came_from_binja: self.send('set_ip', ip=addr - self.bv.start) if self.ip: self.highlight(self.ip, no_color) self.highlight(addr, ip_color) self.ip = addr self.bv.navigate(self.bv.view, addr) # "go to address" equivalent # highlight the target of a branch instruction func = self.bv.get_functions_containing(addr)[0] il_instr = func.get_lifted_il_at(addr) if il_instr.operation != LowLevelILOperation.LLIL_IF or not regs: return instr, jump = self.bv.get_disassembly(addr).split() if ((instr in ('jo') and regs['of'] == 1) or (instr in ('jno') and regs['of'] == 0) or (instr in ('js') and regs['sf'] == 1) or (instr in ('jns') and regs['sf'] == 0) or (instr in ('je', 'jz') and regs['zf'] == 1) or (instr in ('jne') and regs['zf'] == 0) or (instr in ('jb', 'jnae', 'jc') and regs['cf'] == 1) or (instr in ('jnb', 'jae', 'jnc') and regs['cf'] == 0) or (instr in ('jbe', 'jna') and (regs['cf'] == 1 or regs['zf'] == 1)) or (instr in ('ja', 'jnbe') and regs['cf'] == 0 and regs['zf'] == 0) or (instr in ('jl', 'jnge') and regs['sf'] != regs['of']) or (instr in ('jge', 'jnl') and regs['sf'] == regs['of']) or (instr in ('jle', 'jng') and (regs['zf'] == 1 or regs['sf'] != regs['of'])) or (instr in ('jg', 'jnle') and regs['zf'] == 0 and regs['sf'] == 0) or (instr in ('jp', 'jpe') and regs['pf'] == 1) or (instr in ('jnp', 'jpo') and regs['pf'] == 1) or (instr in ('jcxz', 'jecxz') and regs['cx'] == 0)): target = int(jump, 16) else: target = addr + self.bv.get_instruction_length(addr) self.highlight(target, cond_jump_color) def set_bp(self, addr, came_from_binja=False): if came_from_binja: self.send('set_bp', addr=addr - self.bv.start) self.bps.add(addr) self.highlight(addr, enabled_bp_color) def delete_bp(self, addr, came_from_binja=False): if addr not in self.bps: return if came_from_binja: self.send('delete_bp', addr=addr - self.bv.start) self.bps.remove(addr) self.highlight(addr, no_color) def bp_hit(self, addr, regs=None): self.set_ip(addr, regs=regs) self.highlight(addr, hit_bp_color) def vtable(self, ip, target, object, instr): target_sym = self.bv.get_symbol_at(target) target_sym = target_sym.full_name if target_sym else str(target) if object: object_sym = self.bv.get_symbol_at(object) object_sym = object_sym.full_name.split( "::")[0] if object_sym else str(object) comment = '{} (from {} object)'.format(target_sym, object_sym) if instr == 'lea': comment = '(ptr)' + comment else: comment = target_sym self.bv.get_functions_containing(ip)[0].set_comment_at(ip, comment)
class IdaWorker(object): def __init__(self, conn_addr, *args, **kwargs): super(IdaWorker, self).__init__(*args, **kwargs) self.daemon = True self.conn = Client(conn_addr) self.stop = False self.pytest_config = None from PyQt5.QtWidgets import qApp self.qapp = qApp def run(self): try: while not self.stop: command = self.recv() response = self.handle_command(*command) if response: self.conn.send(response) except RuntimeError: log.exception("Runtime error encountered during message handling") except EOFError: log.info("remote connection closed abruptly, terminating.") def recv(self): while not self.stop: self.qapp.processEvents() if not self.conn.poll(1): continue return self.conn.recv() def send(self, *s): return self.conn.send(s) def handle_command(self, command, *command_args): handler_name = "command_" + command if not hasattr(self, handler_name): raise RuntimeError("Unrecognized command recieved: " "'{}'".format(command)) log.debug("Received command: {} with args {}".format( command, command_args)) response = getattr(self, handler_name)(*command_args) log.debug("Responding: {}".format(response)) return response def command_dependencies(self, action, *plugins): # test pytest is installed and return ready if it is if action == "check": try: import pytest del pytest for plugin in plugins: __import__(plugin) return ('dependencies', 'ready') except ImportError: pass # pytest is missing, we'll report so and expect to be requested to # install self.send('dependencies', 'missing') elif action == "install": # test at least pip exists, otherwise we're doomed to fail try: import pip del pip except ImportError: return ('dependencies', 'failed') # handle different versions of pip try: from pip import main as pip_main except ImportError: # here be dragons from pip._internal import main as pip_main # ignoring installed six and upgrading is requried to avoid an osx # bug see https://github.com/pypa/pip/issues/3165 for more details pip_command = ['install', 'pytest'] + list(plugins) if platform.system() == 'Darwin': pip_command += [ '--upgrade', '--user', '--ignore-installed', 'six' ] pip_main(pip_command) # make sure pytest was successfully installed try: import pytest del pytest return ('dependencies', 'ready') except ImportError: return ('dependencies', 'failed') else: raise RuntimeError("Unexpected dependencies argument: " "{}".format(action)) @staticmethod def command_autoanalysis(action): if action == "wait": ida_auto.auto_wait() return ( 'autoanalysis', 'done', ) else: raise RuntimeError("Invalid action received for command: " "{}".format(action)) def command_configure(self, args, option_dict): from _pytest.config import Config import plugin_worker self.pytest_config = Config.fromdictargs(option_dict, args) self.pytest_config.args = args plugin = plugin_worker.WorkerPlugin(worker=self) self.pytest_config.pluginmanager.register(plugin) return ('configure', 'done') def command_cmdline_main(self): self.send('cmdline_main', 'start') self.pytest_config.hook.pytest_cmdline_main(config=self.pytest_config) self.send('cmdline_main', 'finish') @staticmethod def command_ping(): return ('pong', ) def command_quit(self): self.stop = True return ("quitting", )
class CKnowledgeClient: def __init__(self, addressTuple, authkey=b'secret password', oLog=None): self.addressTuple = addressTuple self.authkey = authkey self.conn = None self.oLog = oLog def connect(self): try: self.conn = Client(self.addressTuple, authkey=self.authkey) except: self.modMsg('connection is refused') return False else: return True def send(self, msg): if (not self.closedFlag): return self.conn.send(msg) else: self.modMsg('connection is closed') return False def recv(self, block=False): if (not self.closedFlag): if block: return self.conn.recv() if (self.conn.poll()): return self.conn.recv() else: return ReturnCode.WARNING_NOTHING_TO_RECV else: self.modMsg('connection is closed') return False def close(self): if (not self.closedFlag): self.conn.send('close') try: msg = self.conn.recv() print(msg) if (msg == 'All is closed'): self.conn.close() return True else: return False except: self.modMsg('nothing to recv') return False else: self.modMsg('connection is closed') return False def modMsg(self, string: str): string = self.__class__.__name__ + ' warning: ' + string if (self.oLog != None): self.oLog.safeRecordTime(string) else: print(string) @property def closedFlag(self): if (self.conn == None): return True else: return self.conn.closed
class VariableMonitoringThread(Thread): def __init__(self, **kwargs): Thread.__init__(self) # If PLC name is set, we monitor all PLC tags self.plc_name = kwargs.get('plc_name', None) self.motor = None # If no PLC name is set, we check if motor tags should be monitored if self.plc_name is None: self.motor = kwargs.get('motor', None) # No motor supplied, building thread failed. Either PLC name (for monitoring all PLC tags) or motor # (for monitoring motor/PLC tags) must be supplied if self.motor is None: raise ValueError("No PLC name or motor supplied.") else: self.plc_name = self.motor.plc.name self.listener_ready = False self.stop_event = threading.Event() def __create_connection(self): tmp_base = get_tmp_base_path_from_mkfile() plc_base_path = os.path.join(tmp_base, self.plc_name) address = os.path.join(plc_base_path, 'plc_socket') # Wait until listener is ready # TODO: Refactoring needed (anti-pattern) while not self.listener_ready: try: self.conn = Client(address) self.listener_ready = True except socket.error as e: # Check if listener is not yet ready if e.errno == ENOENT: # socket.error: [Errno 2] No such file or directory sleep(1) else: logger.exception("Unknown socket error occurred.") def run(self): self.__create_connection() tag_names = [] # If motor is none, all PLC tags should be monitored if self.motor is None: self.conn.send(GetAllTagNamesMessage()) else: tag_names = self.motor.plc_vars_map.keys() self.conn.send(MonitorMessage(tag_names)) while not self.stop_event.is_set(): try: if self.conn.poll(): result = self.conn.recv() if isinstance(result, UnknownPlcTagException): raise UnknownPlcTagException(result) elif isinstance(result, TerminateMessage): logger.info("Terminating monitoring of '{}'.".format( ', '.join(tag_names))) break elif isinstance(result, GetAllTagNamesResponseMessage): tag_names = result.tag_names # Recreate connection before sending new message self.listener_ready = False self.__create_connection() self.conn.send(MonitorMessage(tag_names)) elif isinstance(result, MonitorResponseMessage): # First check if we are currently monitoring a PLC or motor if self.motor is None: dev_name = self.plc_name else: dev_name = self.motor.name # We have to alter the name of the result set, as we are receiving PLC tag # changes, yet we have to transmit the motor tag # The 'plc_vars_map' will contain the mapping: {PLC_TAG_NAME: IDX_MOTOR_VARS} idx_motor_var = self.motor.plc_vars_map.get( result.name) if idx_motor_var is not None: # Set the correct motor tag name result.name = self.motor.vars[idx_motor_var][ 'name'] else: logger.error( "Received unknown PLC tag (motor tag map)." ) if dev_name is not None: state_logger.info( "[%s] '%s' variable changed [%s=%s].", result.timestamp.strftime( '%Y-%m-%d %H:%M:%S.%f')[:-3], dev_name, result.name, result.value) else: logger.error("Received unexpected message type '%s'.", type(result)) except EOFError: logger.exception("Received EOF.") break except UnknownPlcTagException: logger.exception("Unknown PLC Tag.") try: # Announce to PLC supervisor that we stop monitoring PLC tags logger.info("Announcing PLC supervisor to stop monitoring.") self.conn.send(StopMonitoringMessage()) except IOError, e: if e.errno == errno.EPIPE: pass else: logger.exception( "Could not announce to PLC supervisor to stop monitoring states for state logging." ) self.conn.close() logger.info("Terminated VariableMonitoringThread.")