def disable_bot(self, msg): """ This will set the enabled-flag to False. :rtype : (bool, str) :type msg: hystck.net.proto.genericmessage_pb2.GenericMessage :param msg: The message that was received (not used by this method) """ self.enabled = False self.logger.info("agent connector is disabled") # print 'agent connector is disabled' m = messagegenerator.generate_announce_message( self.message_id, announcemessage_pb2.ONLINE_DISABLED, self.delegate_class) msg = netutility.NetUtility.length_prefix_message( m.SerializeToString()) self.client_socket.send(msg) self.message_id += 1 return True, "success"
def stop(self): """ Stops the AgentConnector. """ self.logger.info("shutting down the agent connector") # print 'shutting down the agent connector' shutdown_msg = messagegenerator.generate_announce_message( self.message_id, announcemessage_pb2.GRACEFUL_SHUTDOWN, self.delegate_class) try: self.client_socket.send( netutility.NetUtility.length_prefix_message( shutdown_msg.SerializeToString())) except socket.error: pass self.client_close_connection() self.stop_command_listener() self.logger.info("agent connector stopped")
def terminate_bot(self, msg): """ This will terminate the process and send a notification that this was intended to the agent. :rtype : None :type msg: hystck.net.proto.genericmessage_pb2.GenericMessage :param msg: The message that was received (not used by this method) """ self.logger.info("going to terminate process") # print 'going to terminate process' # Tell the agent that this was planned. shutdown_msg = messagegenerator.generate_announce_message( self.message_id, announcemessage_pb2.GRACEFUL_SHUTDOWN, self.delegate_class) try: self.client_socket.send( netutility.NetUtility.length_prefix_message( shutdown_msg.SerializeToString())) except socket.error: pass sys.exit(0)
def listen_thread(self): """ This thread will listen to incoming commands by an agent instance and complete the request. Upon start it will tell the agent instance that it's ready by sending a state announcement. """ if self.enabled: hello_msg = messagegenerator.generate_announce_message( self.message_id, announcemessage_pb2.ONLINE_ENABLED, self.delegate_class) else: hello_msg = messagegenerator.generate_announce_message( self.message_id, announcemessage_pb2.ONLINE_DISABLED, self.delegate_class) self.message_id += 1 # tell agent that bot is ready self.client_socket.send( netutility.NetUtility.length_prefix_message( hello_msg.SerializeToString())) # request the globals globals_request_msg = messagegenerator.generate_globals_request_message( self.message_id) self.client_socket.send( netutility.NetUtility.length_prefix_message( globals_request_msg.SerializeToString())) self.message_id += 1 # start processing messages while self.active: try: # receive and de-marshall the message msg = netutility.NetUtility.receive_prefixed_message( self.client_socket) m = genericmessage_pb2.GenericMessage() m.ParseFromString(msg) self.logger.debug("got message: %s", str(m)) # print "agent connector: got message>\n", str(m) f = self.id_mapper[m.message_type] # get the mapped function ret_val = f( m) # call and pass the message to it's mapped handler if ret_val is None: # if ret_val is not set, default to success ret_val = (True, "success") # generate an answer ret_msg = messagegenerator.generate_answer_message( self.message_id, ret_val[0], msg, ret_val[1]) try: # send the answer back self.client_socket.send( netutility.NetUtility.length_prefix_message( ret_msg.SerializeToString())) except socket.error as e: self.logger.error("Failure sending answer message: %s", e) except KeyError as e: # expect this if a unmapped message type was received self.logger.error("KeyError: %s", e) # print e except socket.error as e: # expect this if a serious problem has occurred with the socket self.logger.error("Socket Error: %s", e) # print "Socket Error: ", str(e) self.client_socket.close() break except RuntimeError as e: # expect this if the reading end or the socket itself has been closed self.logger.warning("Runtime Error: %s", e) # print e.message if e.message == 'unexpected connection close': try: self.client_socket.shutdown(socket.SHUT_RDWR) except socket.error: pass finally: self.client_socket.close() self.logger.info("socket has been closed on demand") # print 'socket has been closed by demand' break else: self.logger.error("RuntimeError: %s", e) self.client_socket.close() break except message.Error as e: # a unhandled error happened self.logger.error("ProtoBuf Error: %s", e)
def agent_server_processor_thread(self, sock): """ This is the bot side message listener thread. :type sock: socket._socketobject :param sock: the socket to work with """ is_graceful = False while self.active: # msg = '' try: # receive a message msg = NetUtility.receive_prefixed_message(sock) m = GenericMessage() m.ParseFromString(msg) self.logger.debug("(bot-side): got message> %s", str(m)) # print "agent (bot-side): got message>\n", str(m) # check if bot wants to go down if m.message_type == messagetypes_pb2.ANNOUNCE: if m.HasExtension(announcemessage_pb2.status_info): if m.Extensions[ announcemessage_pb2. status_info].state == announcemessage_pb2.GRACEFUL_SHUTDOWN: is_graceful = True self.logger.info("expecting graceful shutdown") # print "expecting a graceful shutdown" # pass message to monitor self.client_socket.send(NetUtility.length_prefix_message(msg)) except socket.error as e: # except this to happen if socket was close immaturely self.logger.warning("Socket Error: %s", e) # print "Socket Error: ", str(e) # check if we got a shutdown announcement if not is_graceful: try: # the bot crashed most likely, tell the monitor self.logger.error("Instance has crashed!") cm = messagegenerator.generate_announce_message( 0, announcemessage_pb2.CRASHED, 'AgentConnectorBaseDelegate') self.client_socket.send( NetUtility.length_prefix_message( cm.SerializeToString())) except socket.error: pass break # close the socket anyway since it is faulty self.logger.warning("Closing socket!") sock.close() self.socket_pool.remove(sock) break except RuntimeError as e: # except this to happen if the bot wants to close the socket if e.message == 'unexpected connection close': try: sock.shutdown(socket.SHUT_RDWR) except socket.error: pass finally: sock.close() self.logger.info("socket has been closed by demand") # print 'socket has been closed by demand' self.socket_pool.remove(sock) # check if we got a shutdown announcement if not is_graceful: try: # the bot crash most likely, tell the monitor self.logger.error("Instance has crashed!") cm = messagegenerator.generate_announce_message( 0, announcemessage_pb2.CRASHED, 'AgentConnectorBaseDelegate') self.client_socket.send( NetUtility.length_prefix_message( cm.SerializeToString())) except socket.error: pass break else: self.logger.error("RuntimeError: %s", e) sock.close() self.socket_pool.remove(sock) break except message.Error as e: self.logger.error("ProtoBuf Error: %s", e)