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.")
