def proccess_input(self, clientsock):
""""
Continuously read and process commands from our socket
"""
COMMANDS = {
G.CTRL_CMD_START: self.command_start,
G.CTRL_CMD_STOP: self.command_analysis,
G.CTRL_CMD_PAUSE: self.command_analysis,
G.CTRL_CMD_UNPAUSE: self.command_analysis,
G.CTRL_CMD_LIST: self.command_list,
G.CTRL_CMD_PICKLE: self.command_pickle
}
try:
# Loop forever
while self.RUNNING:
# Get the data
try:
data = G.read_socket_data(clientsock)
except socket.timeout:
# self._cleanup_pointers()
continue
if not data:
break
# cleanup our pointers to give a
# self._cleanup_pointers()
# Split up our command
cmd = LophiCommand.from_data(data) #.rstrip().split(" ")
logger.debug("Got command: %s" % cmd)
# See if it's valid command
if cmd.cmd == "help":
G.send_socket_data(
clientsock,
"The available commands are: %s\n" % (COMMANDS.keys()))
elif cmd.cmd not in COMMANDS.keys():
G.send_socket_data(clientsock,
"Invalid Command: %s\n" % cmd.cmd)
else:
rtn = COMMANDS[cmd.cmd](cmd, clientsock)
if rtn == True:
G.send_socket_data(clientsock, "Success.")
elif rtn == False:
G.send_socket_data(clientsock, "Failure.")
except socket.error:
logger.warning("Looks like our socket closed...")
def proccess_input(self, clientsock):
""""
Continuously read and process commands from our socket
"""
COMMANDS = {G.CTRL_CMD_START:self.command_start,
G.CTRL_CMD_STOP:self.command_analysis,
G.CTRL_CMD_PAUSE:self.command_analysis,
G.CTRL_CMD_UNPAUSE:self.command_analysis,
G.CTRL_CMD_LIST:self.command_list,
G.CTRL_CMD_PICKLE:self.command_pickle}
try:
# Loop forever
while self.RUNNING:
# Get the data
try:
data = G.read_socket_data(clientsock)
except socket.timeout:
# self._cleanup_pointers()
continue
if not data:
break
# cleanup our pointers to give a
# self._cleanup_pointers()
# Split up our command
cmd = LophiCommand.from_data(data) #.rstrip().split(" ")
logger.debug("Got command: %s" % cmd)
# See if it's valid command
if cmd.cmd == "help":
G.send_socket_data(clientsock, "The available commands are: %s\n" %
(COMMANDS.keys()))
elif cmd.cmd not in COMMANDS.keys():
G.send_socket_data(clientsock,"Invalid Command: %s\n" % cmd.cmd)
else:
rtn = COMMANDS[cmd.cmd](cmd, clientsock)
if rtn == True:
G.send_socket_data(clientsock, "Success.")
elif rtn == False:
G.send_socket_data(clientsock, "Failure.")
except socket.error:
logger.warning("Looks like our socket closed...")
def get_analysis(self):
""" Get protocol buffer version of remote analysis """
while 1:
try:
logger.debug("Getting analysis list for Controller/%s" %
self.name)
# Get reply
cmd = LophiCommand(G.CTRL_CMD_PICKLE, args=["analysis"])
analysis_buf = self.send_cmd(cmd)
status = G.read_socket_data(self.SOCK)
# unpack protocol buffer
self.analysis = ProtoBuf.unpack_analysis_list(analysis_buf)
return status
except:
self.connect()
def get_machines(self):
""" Get protocol buffer version of remote machines """
while 1:
try:
logger.debug("Getting machine list for Controller/%s" %
self.name)
# Get response
cmd = LophiCommand(G.CTRL_CMD_PICKLE, args=["machines"])
data = self.send_cmd(cmd)
status = G.read_socket_data(self.SOCK)
# Unpack our machine list
# (WARNING: This a subset of the objects at the server
if data is not None:
self.machines = ProtoBuf.unpack_machine_list(data)
else:
self.machines = []
return status
except:
G.print_traceback()
self.connect()
def start(self):
"""
Main function to just loop forever while waiting for input over amqp
"""
quit_commands = ['q', 'quit', 'exit']
# Setup our handler to close gracefully
G.set_exit_handler(self.cleanup)
# Setup or queues
self.manager = multiprocessing.Manager()
self.INPUT_QUEUE = self.manager.Queue()
# set of comm processes (rabbitmq, etc.) - for cleanup later
self.comm_processes = set([])
# Set up communication queue with all of our consumers, processes, and producers
self.in_queue = multiprocessing.Queue()
self.out_queue = multiprocessing.Queue()
# Listen for physical cards registering
# HOST = ''
# PORT = G.CARD_REG_PORT
#
# self.reg_consumer = Card_Reg_Server((HOST, PORT), UDPHandler, self.in_queue)
# self.reg_consumer.start()
# self.comm_processes.add(self.reg_consumer)
# Setup RabbitMQ consumers and queues
logger.debug("Starting up LOPHI RabbitmQ Producers...")
# self.ctrl_producer = rabbitmq.LOPHI_RPC_Producer(self.amqp_host,
# self.out_queue,
# G.RabbitMQ.CTRL_OUT,
# G.RabbitMQ.CTRL_IN,
# exchange_type=G.RabbitMQ.TYPE_FANOUT,
# exchange=G.RabbitMQ.EXCHANGE_FANOUT)
self.ctrl_producer = rabbitmq.LOPHI_RabbitMQ_Producer(
self.amqp_host,
self.out_queue,
G.RabbitMQ.CTRL_OUT,
exchange_type=G.RabbitMQ.TYPE_FANOUT,
routing_key='',
exchange=G.RabbitMQ.EXCHANGE_FANOUT)
self.ctrl_producer.start()
self.comm_processes.add(self.ctrl_producer)
logger.debug("Starting up LOPHI RabbitMQ Consumers...")
# Listen for control messages, e.g. from a CLI
# self.ctrl_consumer = rabbitmq.LOPHI_RPC_Consumer(self.amqp_host,
# self.in_queue,
# G.RabbitMQ.CTRL_IN)
self.ctrl_consumer = rabbitmq.LOPHI_RabbitMQ_Consumer(
self.amqp_host, self.in_queue, G.RabbitMQ.CTRL_IN)
self.ctrl_consumer.start()
self.comm_processes.add(self.ctrl_consumer)
# Connect to all of our controllers
for c in self.config_list:
self.config_list[c].connect()
print "Waiting for input from queues."
# Just loop forever taking input from rabbitmq
while 1:
user_input = self.in_queue.get()
# Decode input from rabbitmq format
try:
# (corr_id, routing_key, msg) = json.loads(user_input)
# type is type of message
# command
# (type, cmd_data) = msg
cmd = LophiCommand.from_data(user_input)
except:
print "Unknown command: ", user_input
continue
logger.debug("Received msg %s" % cmd)
# check if type is valid
# if msg.type not in self.MSG_TYPES:
# print "Invalid message type: %s\n" % type
# See if it's valid command
if cmd.cmd not in self.COMMANDS.keys():
resp = "Invalid Command: %s\n" % cmd.cmd
logger.debug(resp)
response = cmd.make_response(resp)
self.out_queue.put(response)
else:
# self.process_cmd(type, cmd, corr_id, routing_key)
self.process_cmd(cmd)
"""
@todo: add command to kill master
"""
# Call our cleanup function and shutdown nicely
self.cleanup(None)
def start(self):
"""
Main function to just loop forever while waiting for input over amqp
"""
quit_commands = ['q', 'quit', 'exit']
# Setup our handler to close gracefully
G.set_exit_handler(self.cleanup)
# Setup or queues
self.manager = multiprocessing.Manager()
self.INPUT_QUEUE = self.manager.Queue()
# set of comm processes (rabbitmq, etc.) - for cleanup later
self.comm_processes = set([])
# Set up communication queue with all of our consumers, processes, and producers
self.in_queue = multiprocessing.Queue()
self.out_queue = multiprocessing.Queue()
# Listen for physical cards registering
# HOST = ''
# PORT = G.CARD_REG_PORT
#
# self.reg_consumer = Card_Reg_Server((HOST, PORT), UDPHandler, self.in_queue)
# self.reg_consumer.start()
# self.comm_processes.add(self.reg_consumer)
# Setup RabbitMQ consumers and queues
logger.debug("Starting up LOPHI RabbitmQ Producers...")
# self.ctrl_producer = rabbitmq.LOPHI_RPC_Producer(self.amqp_host,
# self.out_queue,
# G.RabbitMQ.CTRL_OUT,
# G.RabbitMQ.CTRL_IN,
# exchange_type=G.RabbitMQ.TYPE_FANOUT,
# exchange=G.RabbitMQ.EXCHANGE_FANOUT)
self.ctrl_producer = rabbitmq.LOPHI_RabbitMQ_Producer(self.amqp_host,
self.out_queue,
G.RabbitMQ.CTRL_OUT,
exchange_type=G.RabbitMQ.TYPE_FANOUT,
routing_key='',
exchange=G.RabbitMQ.EXCHANGE_FANOUT)
self.ctrl_producer.start()
self.comm_processes.add(self.ctrl_producer)
logger.debug("Starting up LOPHI RabbitMQ Consumers...")
# Listen for control messages, e.g. from a CLI
# self.ctrl_consumer = rabbitmq.LOPHI_RPC_Consumer(self.amqp_host,
# self.in_queue,
# G.RabbitMQ.CTRL_IN)
self.ctrl_consumer = rabbitmq.LOPHI_RabbitMQ_Consumer(self.amqp_host,
self.in_queue,
G.RabbitMQ.CTRL_IN)
self.ctrl_consumer.start()
self.comm_processes.add(self.ctrl_consumer)
# Connect to all of our controllers
for c in self.config_list:
self.config_list[c].connect()
print "Waiting for input from queues."
# Just loop forever taking input from rabbitmq
while 1:
user_input = self.in_queue.get()
# Decode input from rabbitmq format
try:
# (corr_id, routing_key, msg) = json.loads(user_input)
# type is type of message
# command
# (type, cmd_data) = msg
cmd = LophiCommand.from_data(user_input)
except:
print "Unknown command: ", user_input
continue
logger.debug("Received msg %s" % cmd)
# check if type is valid
# if msg.type not in self.MSG_TYPES:
# print "Invalid message type: %s\n" % type
# See if it's valid command
if cmd.cmd not in self.COMMANDS.keys():
resp = "Invalid Command: %s\n" % cmd.cmd
logger.debug(resp)
response = cmd.make_response(resp)
self.out_queue.put(response)
else:
# self.process_cmd(type, cmd, corr_id, routing_key)
self.process_cmd(cmd)
"""
@todo: add command to kill master
"""
# Call our cleanup function and shutdown nicely
self.cleanup(None)
def get_incomplete_analysis(options,positionals):
"""
Resubmit all incomplete analyses
"""
out_queue = multiprocessing.Queue()
ctrl_producer = rabbitmq.LOPHI_RabbitMQ_Producer(options.services_host,
out_queue,
G.RabbitMQ.CTRL_IN)
ctrl_producer.start()
uri = 'mongodb://'+options.services_host+':27017/'
print "* Connecting to %s..."%uri
# Initialize our database connections
client = MongoClient(uri)
# Loop over all of our analyses.
db = client.lophi_db
# Get a list of all of our samples
samples_db = db.samples
analyses = db.analyses
samples = []
for sample_entry in samples_db.find():
has_completed = False
cmd = None
for analysis in analyses.find({"sample": sample_entry[
'sample']}):
if analysis['status'] == "COMPLETED":
has_completed = True
else:
if analysis['machine_type'] == "":
analysis['machine_type'] = options.machine_type
cmd = LophiCommand(cmd=G.CTRL_CMD_START,
analysis=analysis['analysis_script'],
machine_type=analysis['machine_type'],
machine=None,
volatility_profile=analysis['volatility_profile'],
sample_doc_id=analysis['sample'],
submitter=G.get_username_local())
if not has_completed and cmd is not None:
print "* Re-submitting sample (%s)"%sample_entry['sample']
out_queue.put(str(cmd))
# for analysis in analyses.find():
#
#
# if analysis['status'] != "COMPLETED":
#
# print "* Resubmitting %s (Status: %s)" % (analysis['_id'],
# analysis['status'])
#
# if analysis['machine_type'] == "":
# continue
#
# # # Prepare a job to send to the machine, using the sample doc id
# # cmd = LophiCommand(cmd=G.CTRL_CMD_START,
# # analysis=analysis['analysis_script'],
# # machine_type=analysis['machine_type'],
# # machine=None,
# # volatility_profile=analysis['volatility_profile'],
# # sample_doc_id=analysis['sample'],
# # submitter=G.get_username_local())
# #
# # out_queue.put(str(cmd))
time.sleep(5)
out_queue.put(G.CTRL_CMD_KILL)
ctrl_producer.stop()
