def test_client_receives_event_on_topic_only_after_subscribe(self):
"""
The idea of this test is to send an event to a topic which we are not
subscribed, so we shouldn't be notified. Then, we subscribe to that
topic and send a new event, we should get that last one.
"""
with self.create_client() as client:
test_topic = '/test/whatever/' + client.config._client_id
client.connect()
self.assertTrue(client.connected)
# Set request callback (use mock to easily check when it was called)
ecallback = EventCallback()
ecallback.on_event = Mock()
client.add_event_callback(test_topic, ecallback, False)
# Send event thru dxl fabric to a topic which we are *not* subscribed
msg = Event(destination_topic=test_topic)
client.send_event(msg)
time.sleep(1)
# We haven't been notified
self.assertEqual(ecallback.on_event.call_count, 0)
# Subscribe to topic
client.subscribe(test_topic)
# Send event thru dxl fabric again to that topic
msg = Event(destination_topic=test_topic)
client.send_event(msg)
time.sleep(1)
# Now we should have been notified of the event
self.assertEqual(ecallback.on_event.call_count, 1)
def test_firstinstancecallback(self):
class MyFirstInstanceCallback(FirstInstanceCallback):
def __init__(self):
super(MyFirstInstanceCallback, self).__init__()
self.first_instance_dict_received = None
self.original_event_received = None
def on_first_instance(self, first_instance_dict, original_event):
self.first_instance_dict_received = first_instance_dict
self.original_event_received = original_event
first_instance_event_payload = {
RepChangeEventProp.HASHES: [{
"type": HashType.SHA1,
"value": "LWykUGG3lyMS4A5ZM/3/lbuQths="
}, {
"type": HashType.MD5,
"value": "MdvozEQ9LKf9I2rAClL7Fw=="
}, {
"type":
HashType.SHA256,
"value":
"qjxGHUwho5LjctDWykzrHk2ICY1YdllFTq9Nk8ZhiA8="
}],
DetectionEventProp.SYSTEM_GUID:
"testGuid",
DetectionEventProp.NAME:
"MORPH.EXE"
}
first_instance_expected = {
DetectionEventProp.SYSTEM_GUID: "testGuid",
RepChangeEventProp.HASHES: {
HashType.SHA256:
"aa3c461d4c21a392e372d0d6ca4ceb1e4d88098d587659454eaf4d93c661880f",
HashType.SHA1: "2d6ca45061b7972312e00e5933fdff95bb90b61b",
HashType.MD5: "31dbe8cc443d2ca7fd236ac00a52fb17"
},
DetectionEventProp.NAME: "MORPH.EXE"
}
test_event = Event(TEST_TOPIC)
# Set the payload
test_event.payload = json.dumps(first_instance_event_payload)\
.encode(encoding="UTF-8")
first_instance_callback = MyFirstInstanceCallback()
first_instance_callback.on_event(test_event)
self.assertIn("MORPH.EXE", str(test_event.payload))
self.assertDictEqual(
first_instance_callback.first_instance_dict_received,
first_instance_expected)
self.assertEqual(first_instance_callback.original_event_received,
test_event)
def measure_performance(self, client, with_wildcard, topic_exists):
sub_count = 10000
query_multiplier = 10
topic_prefix = "/topic/" + UuidGenerator.generate_id_as_string() + "/"
event_count = [0]
message_ids = set()
payload = UuidGenerator.generate_id_as_string()
payload_as_bytes = payload.encode()
message_id_condition = Condition()
callback = EventCallback()
def on_event(event):
if event.payload == payload_as_bytes:
with message_id_condition:
event_count[0] += 1
message_ids.add(event.message_id)
message_id_condition.notify()
if len(message_ids) % sub_count == 0:
print("Messages size: " + str(len(message_ids)))
callback.on_event = on_event
client.add_event_callback("#", callback, False)
if with_wildcard:
client.subscribe(topic_prefix + "#")
for i in range(sub_count):
if i % 1000 == 0:
print("Subscribed: " + str(i))
client.subscribe(topic_prefix + str(i))
print("Subscribed.")
start_time = time.time()
for j in range(sub_count * query_multiplier):
evt = Event(topic_prefix +
str(j % sub_count +
(sub_count if not topic_exists else 0)))
evt.payload = payload
client.send_event(evt)
with message_id_condition:
while len(message_ids) != sub_count * query_multiplier \
or event_count[0] != sub_count * query_multiplier * (2 if with_wildcard and topic_exists else 1):
current_event = event_count[0]
message_id_condition.wait(5)
if current_event == event_count[0]:
self.fail("Event wait timeout")
self.assertEqual(sub_count * query_multiplier, len(message_ids))
self.assertEqual(
sub_count * query_multiplier *
(2 if with_wildcard and topic_exists else 1), event_count[0])
return time.time() - start_time
def measure_performance(self, client, with_wildcard, topic_exists):
SUB_COUNT = 10000
QUERY_MULTIPLIER = 10
TOPIC_PREFIX = "/topic/" + UuidGenerator.generate_id_as_string() + "/"
event_count = [0]
message_ids = set()
PAYLOAD = UuidGenerator.generate_id_as_string()
message_id_condition = Condition()
cb = EventCallback()
def on_event(event):
if event.payload == PAYLOAD:
with message_id_condition:
event_count[0] += 1
message_ids.add(event.message_id)
message_id_condition.notify()
if len(message_ids) % SUB_COUNT == 0:
print "Messages size: " + str(len(message_ids))
cb.on_event = on_event
client.add_event_callback("#", cb, False)
if with_wildcard:
client.subscribe(TOPIC_PREFIX + "#")
for i in range(SUB_COUNT):
if i % 1000 == 0:
print "Subscribed: " + str(i)
client.subscribe(TOPIC_PREFIX + str(i))
print "Subscribed."
start_time = time.time()
for j in range(SUB_COUNT * QUERY_MULTIPLIER):
evt = Event(TOPIC_PREFIX +
str(j % SUB_COUNT +
(SUB_COUNT if not topic_exists else 0)))
evt.payload = PAYLOAD
client.send_event(evt)
with message_id_condition:
while len(message_ids) != SUB_COUNT * QUERY_MULTIPLIER \
or event_count[0] != SUB_COUNT * QUERY_MULTIPLIER * (2 if with_wildcard and topic_exists else 1):
current_event = event_count[0]
message_id_condition.wait(5)
if current_event == event_count[0]:
self.fail("Event wait timeout")
self.assertEquals(SUB_COUNT * QUERY_MULTIPLIER, len(message_ids))
self.assertEquals(
SUB_COUNT * QUERY_MULTIPLIER *
(2 if with_wildcard and topic_exists else 1), event_count[0])
return time.time() - start_time
def test_event(self):
source_client_guid = UuidGenerator.generate_id_as_string()
source_broker_guid = UuidGenerator.generate_id_as_string()
source_broker_ids = ["{66000000-0000-0000-0000-000000000001}",
"{66000000-0000-0000-0000-000000000002}",
"{66000000-0000-0000-0000-000000000003}"]
source_client_ids = ["{25000000-0000-0000-0000-000000000001}",
"{25000000-0000-0000-0000-000000000002}",
"{25000000-0000-0000-0000-000000000003}"]
source_payload = "EVENT".encode()
event = Event(destination_topic="")
event._source_client_id = source_client_guid
event._source_broker_id = source_broker_guid
event.broker_ids = source_broker_ids
event.client_ids = source_client_ids
event.payload = source_payload
PP.pprint(vars(event))
message = event._to_bytes()
PP.pprint(message)
result = Message._from_bytes(message)
PP.pprint(vars(result))
self.assertEqual(source_client_guid, result.source_client_id)
self.assertEqual(source_broker_guid, result.source_broker_id)
self.assertEqual(source_broker_ids, result.broker_ids)
self.assertEqual(source_client_ids, result.client_ids)
self.assertEqual(source_payload, result.payload)
self.assertEqual(Message.MESSAGE_TYPE_EVENT, result.message_type)
def test_event(self):
source_client_guid = UuidGenerator.generate_id_as_string()
source_broker_guid = UuidGenerator.generate_id_as_string()
event = Event(destination_topic="")
event._source_client_id = source_client_guid
event._source_broker_id = source_broker_guid
event.broker_ids = ["{66000000-0000-0000-0000-000000000001}",
"{66000000-0000-0000-0000-000000000002}",
"{66000000-0000-0000-0000-000000000003}"]
event.client_ids = ["{25000000-0000-0000-0000-000000000001}",
"{25000000-0000-0000-0000-000000000002}",
"{25000000-0000-0000-0000-000000000003}"]
event.payload = str.encode("EVENT")
pp.pprint(vars(event))
message = event._to_bytes()
pp.pprint(message)
result = Message._from_bytes(message)
pp.pprint(vars(result))
assert result.source_client_id == source_client_guid
assert result.source_broker_id == source_broker_guid
assert result.broker_ids == ["{66000000-0000-0000-0000-000000000001}",
"{66000000-0000-0000-0000-000000000002}",
"{66000000-0000-0000-0000-000000000003}"]
assert result.client_ids == ["{25000000-0000-0000-0000-000000000001}",
"{25000000-0000-0000-0000-000000000002}",
"{25000000-0000-0000-0000-000000000003}"]
assert result.payload == str.encode("EVENT")
assert result.message_type == Message.MESSAGE_TYPE_EVENT
def test_execute_sync_during_callback(self):
event_topic = UuidGenerator.generate_id_as_string()
req_topic = UuidGenerator.generate_id_as_string()
with self.create_client() as client:
client.connect()
# callback
def event_callback(_):
with self.event_received_condition:
self.event_received = True
try:
req = Request(destination_topic=req_topic)
client.sync_request(req)
except Exception as ex: # pylint: disable=broad-except
self.request_exception_message = str(ex)
self.event_received_condition.notify_all()
callback = EventCallback()
callback.on_event = event_callback
client.add_event_callback(event_topic, callback)
event = Event(destination_topic=event_topic)
client.send_event(event)
start = time.time()
with self.event_received_condition:
while (time.time() - start < self.MAX_WAIT) and \
not self.event_received:
self.event_received_condition.wait(self.MAX_WAIT)
self.assertIsNotNone(self.request_exception_message)
self.assertIn("different thread", self.request_exception_message)
def test_incoming_message_threading(self):
max_wait = 30
thread_count = 10
thread_name_condition = threading.Condition()
thread_name = set()
event_topic = UuidGenerator.generate_id_as_string()
with self.create_client(
incoming_message_thread_pool_size=thread_count) as client:
client.connect()
event_callback = EventCallback()
def on_event(_):
with thread_name_condition:
thread_name.add(threading.current_thread())
if len(thread_name) == thread_count:
thread_name_condition.notify_all()
event_callback.on_event = on_event
client.add_event_callback(event_topic, event_callback)
for _ in range(0, 1000):
evt = Event(event_topic)
client.send_event(evt)
start = time.time()
with thread_name_condition:
while (time.time() - start < max_wait) and \
len(thread_name) < thread_count:
thread_name_condition.wait(max_wait)
self.assertEqual(thread_count, len(thread_name))
def test_client_send_event_publishes_message_to_dxl_fabric(self):
self.client._client.publish = Mock(return_value=None)
# Create and process Request
msg = Event(destination_topic="")
self.client.send_event(msg)
# Check that callback was called
self.assertEqual(self.client._client.publish.call_count, 1)
def test_execute_sync_during_callback(self):
event_topic = UuidGenerator.generate_id_as_string()
req_topic = UuidGenerator.generate_id_as_string()
with self.create_client() as client:
client.connect()
# callback
def event_callback(event):
try:
req = Request(destination_topic=req_topic)
client.sync_request(req)
except Exception, e:
self.exceptions.append(e)
raise e
ec = EventCallback()
ec.on_event = event_callback
client.add_event_callback(event_topic, ec)
time.sleep(self.SLEEP_TIME) # Check the time
evt = Event(destination_topic=event_topic)
client.send_event(evt)
time.sleep(self.SLEEP_TIME)
self.assertTrue(self.exceptions[0] is not None)
self.assertTrue("different thread" in self.exceptions[0].message)
def test_client_handle_message_with_event_calls_event_callback(self):
event_callback = EventCallback()
event_callback.on_event = Mock()
self.client.add_event_callback(self.test_channel, event_callback)
# Create and process Event
evt = Event(destination_topic=self.test_channel)._to_bytes()
self.client._handle_message(self.test_channel, evt)
# Check that callback was called
self.assertEqual(event_callback.on_event.call_count, 1)
def test_execute_message_other_fields(self):
with self.create_client(max_retries=0) as client:
client.connect()
topic = UuidGenerator.generate_id_as_string()
def on_event(event):
with self.event_received_condition:
try:
self.event_received = event
except Exception as ex: # pylint: disable=broad-except
print(ex)
self.event_received_condition.notify_all()
event_callback = EventCallback()
event_callback.on_event = on_event
client.add_event_callback(topic, event_callback)
event = Event(destination_topic=topic)
event.other_fields = {"key" + str(i): "value" + str(i)
for i in range(self.OTHER_FIELDS_COUNT)}
event.other_fields[b"key_as_bytes"] = b"val_as_bytes"
client.send_event(event)
# Bytes values for other field keys/values are expected to be
# converted to unicode strings as received from the DXL fabric.
del event.other_fields[b"key_as_bytes"]
event.other_fields[u"key_as_bytes"] = u"val_as_bytes"
start = time.time()
with self.event_received_condition:
while (time.time() - start < self.MAX_WAIT) and \
not self.event_received:
self.event_received_condition.wait(self.MAX_WAIT)
self.assertIsNotNone(self.event_received)
self.assertIsNotNone(self.event_received.other_fields)
for i in range(self.OTHER_FIELDS_COUNT):
self.assertEqual(
event.other_fields["key" + str(i)],
self.event_received.other_fields.get("key" + str(i), ""))
self.assertEqual(event.other_fields["key_as_bytes"],
self.event_received.other_fields.get(
"key_as_bytes", ""))
def test_execute_events(self):
with self.create_client(max_retries=0) as client:
try:
client.connect()
topic = UuidGenerator.generate_id_as_string()
# Create and register an event callback. Ensure that all sent events are received
# (via the outstanding events set). Also, track the number of total events received.
def event_callback(event):
with self.event_condition:
# Increment count of responses received
self.event_count += 1
# Remove from outstanding events
self.remove_outstanding_event(event.message_id)
# Notify that a response has been received (are we done yet?)
self.event_condition.notify_all()
callback = EventCallback()
callback.on_event = event_callback
client.add_event_callback(topic, callback)
for _ in range(0, self.EVENT_COUNT):
event = Event(topic)
self.append_outstanding_event(event.message_id)
client.send_event(event)
with self.event_condition:
while self.event_count != self.EVENT_COUNT:
current_count = self.event_count
self.event_condition.wait(self.MAX_EVENT_WAIT)
if current_count == self.event_count:
self.fail("Event wait timeout.")
self.assertEqual(0, len(self.outstanding_events))
print("Events test: PASSED")
except Exception as ex:
print(ex)
raise ex
def test_incoming_message_threading(self):
thread_count = 10
thread_name = set()
event_topic = UuidGenerator.generate_id_as_string()
with self.create_client(incoming_message_thread_pool_size=thread_count) as client:
client.connect()
event_callback = EventCallback()
def on_event(event):
thread_name.add(threading.current_thread())
event_callback.on_event = on_event
client.add_event_callback(event_topic, event_callback)
for i in range(0, 1000):
evt = Event(event_topic)
client.send_event(evt)
time.sleep(30)
self.assertEquals(thread_count, len(thread_name))
def test_event_with_empty_broker_and_client_guids(self):
source_client_guid = UuidGenerator.generate_id_as_string()
source_broker_guid = UuidGenerator.generate_id_as_string()
event = Event(destination_topic="")
event._source_client_id = source_client_guid
event._source_broker_id = source_broker_guid
event.payload = "EVENT".encode()
PP.pprint(vars(event))
message = event._to_bytes()
PP.pprint(message)
result = Message._from_bytes(message)
PP.pprint(vars(result))
self.assertTrue(isinstance(result.broker_ids, list))
self.assertTrue(isinstance(result.client_ids, list))
def run():
try:
with client_factory(max_retries=0) as client:
retries = self.MAX_CONNECT_RETRIES
connected = False
while not connected and retries > 0:
try:
self.connect_time_start = time.time()
client.connect()
connected = True
except Exception: # pylint: disable=broad-except
if retries > 0:
retries -= 1
self.connect_retries += 1
self.assertTrue(connected,
"Unable to connect after retries")
def on_event(event):
with self.event_count_condition:
self.event_count += 1
current_count = self.event_count
if current_count == \
self.EVENT_COUNT * self.THREAD_COUNT:
self.event_count_condition.notify_all()
if current_count % 100 == 0:
print(client.config._client_id + " : " +
str(current_count) + " : " +
event.payload.decode("utf8"))
# callback registration
callback = EventCallback()
callback.on_event = on_event
client.add_event_callback(event_topic, callback)
# Waiting all clients have connected
with self.connect_condition:
self.atomic_connect_count += 1
if self.atomic_connect_count == self.THREAD_COUNT:
self.connect_condition.notify_all()
time_remaining = self.MAX_TIME
while self.atomic_connect_count != \
self.THREAD_COUNT and time_remaining > 0:
self.connect_condition.wait(
timeout=time_remaining)
time_remaining = start - time.time() + \
self.MAX_TIME
self.assertEqual(
self.THREAD_COUNT, self.atomic_connect_count,
"Timeout waiting for all threads to connect")
# Once all clients have connected, reset timing information
if self.requests_start_time == 0:
self.requests_start_time = time.time()
self.connect_time = \
self.requests_start_time - \
self.connect_time_start
for i in range(0, self.EVENT_COUNT):
event = Event(event_topic)
if i % 10 == 0:
print("###send: " + str(i))
event.payload = str(i)
send_client.send_event(event)
with self.event_count_condition:
time_remaining = self.MAX_TIME
while self.event_count != \
self.EVENT_COUNT * self.THREAD_COUNT and \
time_remaining > 0:
self.event_count_condition.wait(
timeout=time_remaining)
time_remaining = start - time.time() + \
self.MAX_TIME
self.assertEqual(
self.EVENT_COUNT * self.THREAD_COUNT,
self.event_count,
"Timed out while receiving events")
if self.requests_end_time == 0:
self.requests_end_time = time.time()
except Exception as ex:
logging.error(ex)
raise ex
def run():
try:
with client_factory(max_retries=0) as client:
retries = self.MAX_CONNECT_RETRIES
connected = False
while not connected and retries > 0:
try:
self.connect_time_start = time.time()
client.connect()
connected = True
except Exception:
if retries > 0:
retries -= 1
self.connect_retries += 1
self.assertTrue(connected,
"Unable to connect after retries")
def on_event(event):
with self.event_count_condition:
self.event_count += 1
current_count = self.event_count
self.event_count_condition.notify_all()
if current_count % 100 == 0:
print client.config._client_id + " : " + str(
current_count) + " : " + event.payload
# callback registration
callback = EventCallback()
callback.on_event = on_event
client.add_event_callback(event_topic, callback)
# Waiting all clients have connected
with self.connect_condition:
self.atomic_connect_count += 1
curr_count = self.atomic_connect_count
self.connect_condition.notify_all()
while self.atomic_connect_count != self.THREAD_COUNT:
self.connect_condition.wait(
timeout=self.MAX_CONNECT_WAIT)
if curr_count == self.atomic_connect_count:
self.fail(
"Timeout waiting for all threads to connect"
)
# Once all clients have connected, reset timing information
if self.requests_start_time == 0:
self.requests_start_time = time.time()
self.connect_time = self.requests_start_time - self.connect_time_start
for i in range(0, self.EVENT_COUNT):
event = Event(event_topic)
if i % 10 == 0:
print "###send: " + str(i)
event.payload = str(i)
send_client.send_event(event)
with self.event_count_condition:
while self.event_count != self.EVENT_COUNT * self.THREAD_COUNT:
curr_count = self.event_count
self.event_count_condition.wait(
timeout=self.MAX_CONNECT_WAIT)
if self.event_count == curr_count:
self.fail(
"Timed out while receiving events")
self.event_count_condition.notify_all()
if self.requests_end_time == 0:
self.requests_end_time = time.time()
except Exception, e:
logging.error(e.message)
raise e
def test_repchangecallback(self):
class MyReputationChangeCallback(ReputationChangeCallback):
def __init__(self):
super(MyReputationChangeCallback, self).__init__()
self.rep_change_dict_received = {}
self.original_event_received = None
def on_reputation_change(self, rep_change_dict, original_event):
self.rep_change_dict_received = rep_change_dict
self.original_event_received = original_event
rep_change_event_payload = {
RepChangeEventProp.OLD_REPUTATIONS: {
"reputations": [{
ReputationProp.TRUST_LEVEL: TrustLevel.NOT_SET,
ReputationProp.PROVIDER_ID: FileProvider.ENTERPRISE,
ReputationProp.CREATE_DATE: 1409783001,
ReputationProp.ATTRIBUTES: {
"2098277": "256",
}
}, {
ReputationProp.TRUST_LEVEL: TrustLevel.KNOWN_TRUSTED,
ReputationProp.PROVIDER_ID: FileProvider.GTI,
ReputationProp.CREATE_DATE: 1409783001,
ReputationProp.ATTRIBUTES: {
GtiAttrib.ORIGINAL_RESPONSE: "2139160704"
}
}],
"props": {
"serverTime": 1409851328
}
},
RepChangeEventProp.NEW_REPUTATIONS: {
"reputations": [{
ReputationProp.TRUST_LEVEL: TrustLevel.MOST_LIKELY_TRUSTED,
ReputationProp.PROVIDER_ID: FileProvider.ENTERPRISE,
ReputationProp.CREATE_DATE: 1409783001,
ReputationProp.ATTRIBUTES: {
"2098277": "256",
}
}, {
ReputationProp.TRUST_LEVEL: TrustLevel.KNOWN_TRUSTED,
ReputationProp.PROVIDER_ID: FileProvider.GTI,
ReputationProp.CREATE_DATE: 1409783001,
ReputationProp.ATTRIBUTES: {
GtiAttrib.ORIGINAL_RESPONSE: "2139160704"
}
}],
"props": {
"serverTime": 1409851328
}
},
FileRepChangeEventProp.RELATIONSHIPS: {
"certificate": {
RepChangeEventProp.HASHES: [{
"value": "rB/QkipKKm5XeazdYodHwoOUsLk=",
"type": HashType.SHA1
}],
"publicKeySha1":
"Q139Rw9ydDfHy08Hy6H5ofQnJlY="
}
},
RepChangeEventProp.HASHES: [{
"type": HashType.MD5,
"value": "bQvLG6j1WmwRB8LZ2gPa1w=="
}, {
"type": HashType.SHA1,
"value": "OxbrjQd0H6+3meBW5YuBoInTcqM="
}, {
"type":
HashType.SHA256,
"value":
"yXfKH1ESH+5YzaiIJ6YXOtTx1y2AJihOTE9EMCqWfkA="
}],
RepChangeEventProp.UPDATE_TIME:
1409851328
}
rep_change_expected = {
RepChangeEventProp.NEW_REPUTATIONS: {
FileProvider.GTI: {
ReputationProp.ATTRIBUTES: {
GtiAttrib.ORIGINAL_RESPONSE: "2139160704"
},
ReputationProp.TRUST_LEVEL: TrustLevel.KNOWN_TRUSTED,
ReputationProp.CREATE_DATE: 1409783001,
ReputationProp.PROVIDER_ID: FileProvider.GTI
},
FileProvider.ENTERPRISE: {
ReputationProp.ATTRIBUTES: {
"2098277": "256"
},
ReputationProp.TRUST_LEVEL: TrustLevel.MOST_LIKELY_TRUSTED,
ReputationProp.CREATE_DATE: 1409783001,
ReputationProp.PROVIDER_ID: FileProvider.ENTERPRISE
}
},
FileRepChangeEventProp.RELATIONSHIPS: {
"certificate": {
RepChangeEventProp.HASHES: {
HashType.SHA1:
"ac1fd0922a4a2a6e5779acdd628747c28394b0b9"
},
"publicKeySha1": "435dfd470f727437c7cb4f07cba1f9a1f4272656"
}
},
RepChangeEventProp.HASHES: {
HashType.SHA256:
"c977ca1f51121fee58cda88827a6173ad4f1d72d8026284e4c4f44302a967e40",
HashType.SHA1: "3b16eb8d07741fafb799e056e58b81a089d372a3",
HashType.MD5: "6d0bcb1ba8f55a6c1107c2d9da03dad7"
},
RepChangeEventProp.UPDATE_TIME: 1409851328,
RepChangeEventProp.OLD_REPUTATIONS: {
FileProvider.GTI: {
ReputationProp.ATTRIBUTES: {
GtiAttrib.ORIGINAL_RESPONSE: "2139160704"
},
ReputationProp.TRUST_LEVEL: TrustLevel.KNOWN_TRUSTED,
ReputationProp.CREATE_DATE: 1409783001,
ReputationProp.PROVIDER_ID: FileProvider.GTI
},
FileProvider.ENTERPRISE: {
ReputationProp.ATTRIBUTES: {
"2098277": "256"
},
ReputationProp.TRUST_LEVEL: TrustLevel.NOT_SET,
ReputationProp.CREATE_DATE: 1409783001,
ReputationProp.PROVIDER_ID: FileProvider.ENTERPRISE
}
}
}
test_event = Event(TEST_TOPIC)
# Set the payload
test_event.payload = json.dumps(rep_change_event_payload)\
.encode(encoding="UTF-8")
rep_change_callback = MyReputationChangeCallback()
rep_change_callback.on_event(test_event)
self.assertDictEqual(rep_change_callback.rep_change_dict_received,
rep_change_expected)
self.assertEqual(rep_change_callback.original_event_received,
test_event)
def test_detectioncallback(self):
class MyDetectionCallback(DetectionCallback):
def __init__(self):
super(MyDetectionCallback, self).__init__()
self.detection_dict_received = None
self.original_event_received = None
def on_detection(self, detection_dict, original_event):
self.detection_dict_received = detection_dict
self.original_event_received = original_event
detect_event_payload = {
RepChangeEventProp.HASHES: [{
"value": "CZnbhOFq32TBWnuAOUhLMw==",
"type": HashType.MD5
}, {
"value": "7vZcAfgW1DgH2WrHY5A3h14Fbks=",
"type": HashType.SHA1
}, {
"type":
HashType.SHA256,
"value":
"yXfKH1ESH+5YzaiIJ6YXOtTx1y2AJihOTE9EMCqWfkA="
}],
DetectionEventProp.SYSTEM_GUID:
"{abc5d2c6-e959-11e3-baeb-005056c00009}",
DetectionEventProp.REMEDIATION_ACTION:
5,
DetectionEventProp.LOCAL_REPUTATION:
1,
DetectionEventProp.DETECTION_TIME:
1402617156
}
detect_expected = {
DetectionEventProp.REMEDIATION_ACTION: 5,
DetectionEventProp.SYSTEM_GUID:
u"{abc5d2c6-e959-11e3-baeb-005056c00009}",
RepChangeEventProp.HASHES: {
HashType.SHA256:
"c977ca1f51121fee58cda88827a6173ad4f1d72d8026284e4c4f44302a967e40",
HashType.SHA1: "eef65c01f816d43807d96ac7639037875e056e4b",
HashType.MD5: "0999db84e16adf64c15a7b8039484b33"
},
DetectionEventProp.LOCAL_REPUTATION: 1,
DetectionEventProp.DETECTION_TIME: 1402617156
}
test_event = Event(TEST_TOPIC)
# Set the payload
test_event.payload = json.dumps(detect_event_payload)\
.encode(encoding="UTF-8")
detection_callback = MyDetectionCallback()
detection_callback.on_event(test_event)
self.assertDictEqual(detection_callback.detection_dict_received,
detect_expected)
self.assertEqual(detection_callback.original_event_received,
test_event)
def _mcafee_publish_to_dxl_function(self, event, *args, **kwargs):
"""Function: A function which takes 3 inputs:
mcafee_topic_name: String of the topic name. ie: /mcafee/service/epo/remote/epo1.
mcafee_dxl_payload: The text of the payload to publish to the topic.
mcafee_return_request: Specify whether or not to wait for and return the response.
The function will publish the provided payload to the provided topic.
Indicate whether acknowledgment response should be returned."""
try:
yield StatusMessage("Starting...")
# Get the function parameters:
mcafee_topic_name = kwargs.get("mcafee_topic_name") # text
if not mcafee_topic_name:
yield FunctionError("mcafee_topic_name is required")
mcafee_dxl_payload = kwargs.get("mcafee_dxl_payload") # text
if not mcafee_dxl_payload:
yield FunctionError("mcafee_dxl_payload is required")
mcafee_publish_method = self.get_select_param(
kwargs.get("mcafee_publish_method")
) # select, values: "Event", "Service"
if not mcafee_publish_method:
yield FunctionError("mcafee_publish_method is required")
mcafee_wait_for_response = self.get_select_param(
kwargs.get(
"mcafee_wait_for_response")) # select, values: "Yes", "No"
log.info("mcafee_topic_name: %s", mcafee_topic_name)
log.info("mcafee_dxl_payload: %s", mcafee_dxl_payload)
log.info("mcafee_publish_method: %s", mcafee_publish_method)
log.info("mcafee_wait_for_response: %s", mcafee_wait_for_response)
response = None
# Publish Event
if mcafee_publish_method == "Event":
event = Event(mcafee_topic_name)
event.payload = mcafee_dxl_payload
yield StatusMessage("Publishing Event...")
self.client.send_event(event)
# Invoke Service
else:
req = Request(mcafee_topic_name)
req.payload = mcafee_dxl_payload
yield StatusMessage("Invoking Service...")
if mcafee_wait_for_response == "No":
self.client.async_request(req)
else:
response = Response(
self.client.sync_request(req, timeout=300))
yield StatusMessage("Done...")
r = {
"mcafee_topic_name": mcafee_topic_name,
"mcafee_dxl_payload": mcafee_dxl_payload,
"mcafee_publish_method": mcafee_publish_method,
"mcafee_wait_for_response": mcafee_wait_for_response
}
# Return response from publishing to topic
if response is not None:
r["response"] = vars(response)
yield FunctionResult(r)
else:
yield FunctionResult(r)
except Exception as e:
yield FunctionError(e)
def test_wildcard_services(self):
max_wait = 10
with self.create_client() as client:
# The request message that the service receives
service_request_message = []
# The request message corresponding to the response received by the client
client_response_message_request = []
# The event that we received
client_event_message = []
client_event_message_condition = Condition()
# The payload that the service receives
service_request_message_receive_payload = []
client.connect()
info = ServiceRegistrationInfo(client, "myWildcardService")
meta = {}
# Transform events mapped to "test/#/" to "request/test/..."
meta["EventToRequestTopic"] = "/test/#"
meta["EventToRequestPrefix"] = "/request"
info.metadata = meta
rcb = RequestCallback()
def on_request(request):
print("## Request in service: " + request.destination_topic + ", " + str(request.message_id))
print("## Request in service - payload: " + request.payload)
service_request_message.append(request.message_id)
service_request_message_receive_payload.append(request.payload)
response = Response(request)
response.payload = "Request response - Event payload: " + request.payload
client.send_response(response)
rcb.on_request = on_request
info.add_topic("/request/test/#", rcb)
client.register_service_sync(info, 10)
evt = Event("/test/bar")
rcb = ResponseCallback()
def on_response(response):
# Only handle the response corresponding to the event we sent
if response.request_message_id == evt.message_id:
print("## received_response: " + response.request_message_id + ", " + response.__class__.__name__)
print("## received_response_payload: " + response.payload)
client_response_message_request[0] = response.request_message_id
rcb.on_response = on_response
client.add_response_callback("", rcb)
ecb = EventCallback()
def on_event(event):
print("## received event: " + event.destination_topic + ", " + event.message_id)
with client_event_message_condition:
client_event_message.append(event.message_id)
client_event_message_condition.notify_all()
ecb.on_event = on_event
client.add_event_callback("/test/#", ecb)
# Send our event
print("## Sending event: " + evt.destination_topic + ", " + evt.message_id)
evt.payload = "Unit test payload"
client.send_event(evt)
start = time.time()
with client_event_message_condition:
while (time.time() - start < max_wait) and \
not client_event_message:
client_event_message_condition.wait(max_wait)
# # Make sure the service received the request properly
# self.assertEqual(evt.message_id, service_request_message[0])
# # Make sure the service received the request payload from the event properly
# self.assertEqual(evt.payload, service_request_message_receive_payload[0])
# Make sure the response we received was for the request message
# self.assertEqual(evt.message_id, client_response_message_request[0])
# Make sure we received the correct event
self.assertGreater(len(client_event_message), 0)
self.assertEqual(evt.message_id, client_event_message[0])
def set_external_file_reputation(self, trust_level, hashes, file_type=0, filename="", comment=""):
"""
Sets the "External" reputation (`trust level`) of a specified file (as identified by hashes).
.. note::
**Client Authorization**
The OpenDXL Python client invoking this method must have permission to send messages to the
``/mcafee/event/external/file/report`` topic which is part of the
``TIE Server Set External Reputation`` authorization group.
The following page provides an example of authorizing a Python client to send messages to an
`authorization group`. While the example is based on McAfee Active Response (MAR), the
instructions are the same with the exception of swapping the ``TIE Server Set External Reputation``
`authorization group` in place of ``Active Response Server API``:
`<https://opendxl.github.io/opendxl-client-python/pydoc/marsendauth.html>`_
**Example Usage**
.. code-block:: python
# Set the External reputation (trust level) for file.exe to Known Trusted
tie_client.set_external_file_reputation(
TrustLevel.KNOWN_TRUSTED, {
HashType.MD5: "f2c7bb8acc97f92e987a2d4087d021b1",
HashType.SHA1: "7eb0139d2175739b3ccb0d1110067820be6abd29",
HashType.SHA256: "142e1d688ef0568370c37187fd9f2351d7ddeda574f8bfa9b0fa4ef42db85aa2"
},
FileType.PEEXE,
filename="notepad.exe",
comment="Reputation set via OpenDXL")
:param trust_level: The new `trust level` for the file. The list of standard `trust levels` can be found in the
:class:`dxltieclient.constants.TrustLevel` constants class.
:param hashes: A ``dict`` (dictionary) of hashes that identify the file to update the reputation for.
The ``key`` in the dictionary is the `hash type` and the ``value`` is the `hex` representation of the
hash value. See the :class:`dxltieclient.constants.HashType` class for the list of `hash type`
constants.
:param file_type: A number that represents the file type. The list of allowed `file types` can be found in the
:class:`dxltieclient.constants.FileType` constants class. (optional)
:param filename: A file name to associate with the file (optional)
:param comment: A comment to associate with the file (optional)
"""
if not trust_level:
raise ValueError("TrustLevel was not specified")
if not self.valid_parameter(FileType, file_type):
raise ValueError("FileType was not a valid entry")
if not self.valid_parameter(TrustLevel, trust_level):
raise ValueError("TrustLevel was not a valid entry")
if not hashes:
raise ValueError("File hashes were not specified")
# Create the event
event = Event(TIE_EVENT_EXTERNAL_FILE_REPORT_TOPIC)
# Create a dictionary for the payload
payload_dict = {
"file": {
"type": file_type,
"hashes": hashes,
"attributes": {
"filename": filename
},
"reputation": {
"score": trust_level
}
},
"provider": {
"id": FileProvider.EXTERNAL
},
"comment": comment
}
# Set the payload
event.payload = json.dumps(payload_dict)
# Send the event
self._dxl_client.send_event(event)
