def _many_message(self, count, drop_count, command):
if type(command) == list:
in_command = os.linesep.join(command)
count = len(command)
else:
in_command = ""
for i in range(count):
in_command = in_command + command + os.linesep
inlines = io.StringIO(in_command)
outfile = io.StringIO()
conn = self._get_conn(inlines, outfile, drop_count)
request_listener = reply.RequestListener(self.conf_obj, conn,
self.disp, self.db)
conn.set_receiver(request_listener)
self.disp.start_workers(request_listener)
# wait until the request is done
while request_listener.get_messages_processed() < count:
dcm_events.poll()
for line in outfile.getvalue().split('\n'):
line = line.strip()
if line:
output = json.loads(line.strip())
self.assertEquals(0, output['return_code'])
request_listener.shutdown()
request_listener.wait_for_all_nicely()
def setUp(self):
self.test_conf_path = \
os.path.join(self.test_base_path, "etc", "agent.conf")
self.conf_obj = config.AgentConfig([self.test_conf_path])
self.svc = service.DCMAgent(self.conf_obj)
self._event = threading.Event()
utils.verify_config_file(self.conf_obj)
# script_dir must be forced to None so that we get the built in dir
self.svc.pre_threads()
self.conf_obj.start_job_runner()
self.disp = dispatcher.Dispatcher(self.conf_obj)
self.req_conn = test_conn.RequestConnection()
self.reply_conn = test_conn.ReplyConnection()
self.db = persistence.SQLiteAgentDB(
os.path.join(self.test_base_path, "etc", "agentdb.sql"))
self.request_listener = reply.RequestListener(self.conf_obj,
self.reply_conn,
self.disp, self.db)
observers = self.request_listener.get_reply_observers()
observers.append(self)
self.req_conn.set_request_listener(self.request_listener)
self.agent_id = "theAgentID" + str(uuid.uuid4())
self.customer_id = 50
handshake_doc = {}
handshake_doc["version"] = "1"
handshake_doc["agentID"] = self.agent_id
handshake_doc["cloudId"] = "Amazon"
handshake_doc["customerId"] = self.customer_id
handshake_doc["regionId"] = None
handshake_doc["zoneId"] = "rack2"
handshake_doc["serverId"] = "thisServer"
handshake_doc["serverName"] = "dcm.testagent.com"
handshake_doc["ephemeralFileSystem"] = "/tmp"
handshake_doc["encryptedEphemeralFsKey"] = "DEADBEAF"
self.svc.conn = self.reply_conn
self.svc.disp = self.disp
self.svc.request_listener = self.request_listener
self.svc.handshaker.incoming_document({
"handshake": handshake_doc,
"return_code": 200
})
self.disp.start_workers(self.request_listener)
def _many_message(self, drop_count, command, retrans_list):
class TestStateObserver(reply.ReplyObserverInterface):
def new_message(self, reply):
pass
def message_done(self, reply):
self.state_change_list = reply._sm.get_event_list()
def incoming_message(self, incoming_doc):
pass
request_listener = None
disp = dispatcher.Dispatcher(self.conf_obj)
try:
in_command = os.linesep.join(command)
count = len(command)
inlines = io.StringIO(in_command)
outfile = io.StringIO()
conn = self._get_conn(inlines, outfile, drop_count, retrans_list)
request_listener = reply.RequestListener(self.conf_obj, conn, disp,
self.db)
conn.set_receiver(request_listener)
to = TestStateObserver()
rol = request_listener.get_reply_observers()
rol.insert(0, to)
disp.start_workers(request_listener)
# wait until the request is done. in the case of reply
# retransmissions this value could be greater than count
while request_listener.get_messages_processed() < count:
dcm_events.poll()
for line in outfile.getvalue().split('\n'):
line = line.strip()
if line:
output = json.loads(line.strip())
self.assertEquals(0, output['return_code'])
finally:
if request_listener:
request_listener.shutdown()
request_listener.wait_for_all_nicely()
if disp:
disp.stop()
return to.state_change_list
def test_read_request_retrans_request(self):
disp = mock.Mock()
conn = mock.Mock()
request_id = "requestID"
message_id = "messageID"
request_doc = {
'type': types.MessageTypes.REQUEST,
'request_id': request_id,
'message_id': message_id,
'payload': {}
}
# just make sure that everything doesnt blow up on a repeated request
reply_listener = reply.RequestListener(self.conf, conn, disp, self.db)
reply_listener.incoming_parent_q_message(request_doc)
reply_listener.incoming_parent_q_message(request_doc)
def test_read_request(self):
conn = mock.Mock()
disp = mock.Mock()
request_id = "requestID"
message_id = "messageID"
request_doc = {
'type': types.MessageTypes.REQUEST,
'request_id': request_id,
'message_id': message_id,
'payload': {}
}
disp.incoming_request.return_value = None
reply_listener = reply.RequestListener(self.conf, conn, disp, self.db)
reply_listener.incoming_parent_q_message(request_doc)
self.assertTrue(disp.incoming_request.called)
def test_unknown_ack(self):
disp = mock.Mock()
conn = mock.Mock()
request_id = "requestID"
message_id = "messageID"
ack_doc = {
"type": types.MessageTypes.ACK,
"request_id": request_id,
"message_id": message_id,
"payload": {}
}
reply_listener = reply.RequestListener(self.conf, conn, disp, self.db)
reply_listener.incoming_parent_q_message(ack_doc)
(param_list, keywords) = conn.send.call_args
send_doc = param_list[0]
self.assertTrue('type' in send_doc)
self.assertEqual(send_doc['type'], types.MessageTypes.NACK)
def _simple_message(self, drop_count, command, stdout, stderr):
inlines = io.StringIO(command)
outfile = io.StringIO()
conn = self._get_conn(inlines, outfile, drop_count)
request_listener = reply.RequestListener(self.conf_obj, conn,
self.disp, self.db)
try:
conn.set_receiver(request_listener)
self.disp.start_workers(request_listener)
# wait until the request is done
while request_listener.get_messages_processed() != 1:
dcm_events.poll()
line = outfile.getvalue().split('\n')[0]
line = line.strip()
output = json.loads(line)
self.assertEquals(stdout, output['message'].strip())
self.assertEquals(stderr, output['error_message'])
self.assertEquals(0, output['return_code'])
finally:
request_listener.shutdown()
request_listener.wait_for_all_nicely()
def run_agent(self):
try:
self.db_cleaner = persistence.DBCleaner(
self._db, self.conf.storage_db_timeout, 100000, 3600)
self.db_cleaner.start()
# def get a connection object
self.conn = config.get_connection_object(self.conf)
self.disp = dispatcher.Dispatcher(self.conf)
# this is done in two steps because we only want to fork before the
# threads are created
if self.conf.intrusion_detection_ossec:
self.intrusion_detection =\
ossec.AlertSender(
self.conn, self._db,
max_process_time=self.conf.intrusion_detection_max_process_time,
alert_threshold=self.conf.intrusion_detection_alert_threshold)
self.intrusion_detection.start()
self.request_listener = reply.RequestListener(
self.conf,
self.conn,
self.disp,
self._db,
id_system=self.intrusion_detection)
logger.set_dcm_connection(self.conf, self.conn)
self.conn.connect(self.request_listener.incoming_parent_q_message,
self.handshaker)
self.disp.start_workers(self.request_listener)
rc = self.agent_main_loop()
return rc
finally:
self.cleanup_agent()