def test_pre_post_recv(self):
pre_callback_future = model.Future()
post_callback_future = model.Future()
expected_bytes = []
expected_rounds = [1, 2]
def pre_cb(read_bytes):
with pre_callback_future:
expected_bytes.insert(0, read_bytes)
self.assertGreater(read_bytes, 3)
pre_callback_future.complete(True)
def post_cb(data):
with post_callback_future:
self.assertEqual(expected_bytes.pop(), len(data))
expected_rounds.pop()
if not expected_rounds:
post_callback_future.complete(True)
conn = self.default_client.connect(self.server, self.port)
conn.register_callback(model.EV_RES_PRE_RECV, pre_cb)
conn.register_callback(model.EV_RES_POST_RECV, post_cb)
conn.negotiate()
self.assertTrue(pre_callback_future.result(timeout=2))
self.assertTrue(post_callback_future.result(timeout=2))
def test_pre_deserialize(self):
callback_future = model.Future()
def cb(data):
with callback_future:
self.assertGreater(len(data), 4)
callback_future.complete(True)
self.default_client.register_callback(model.EV_RES_PRE_DESERIALIZE, cb)
conn = self.default_client.connect(self.server)
conn.negotiate()
self.assertTrue(callback_future.result(timeout=2))
def test_pre_post_send(self):
pre_callback_future = model.Future()
post_callback_future = model.Future()
expected_bytes = []
def pre_cb(data):
with pre_callback_future:
expected_bytes.insert(0, len(data))
self.assertGreater(expected_bytes[0], 16)
pre_callback_future.complete(True)
def post_cb(bytes_written):
with post_callback_future:
self.assertEqual(expected_bytes.pop(), bytes_written)
if not expected_bytes:
post_callback_future.complete(True)
self.default_client.register_callback(model.EV_REQ_PRE_SEND, pre_cb)
self.default_client.register_callback(model.EV_REQ_POST_SEND, post_cb)
conn = self.default_client.connect(self.server)
conn.negotiate()
self.assertTrue(pre_callback_future.result(timeout=2))
self.assertTrue(post_callback_future.result(timeout=2))
def test_pre_serialize(self):
callback_future = model.Future()
def cb(frame):
with callback_future:
self.assertTrue(isinstance(frame, core.Frame))
self.assertTrue(isinstance(frame[0], smb2.Smb2))
self.assertTrue(isinstance(frame[0][0], smb2.NegotiateRequest))
self.assertFalse(hasattr(frame, "buf"))
callback_future.complete(True)
self.default_client.register_callback(model.EV_REQ_PRE_SERIALIZE, cb)
conn = self.default_client.connect(self.server)
conn.negotiate()
self.assertTrue(callback_future.result(timeout=2))
def test_post_deserialize(self):
callback_future = model.Future()
def cb(frame):
with callback_future:
self.assertTrue(isinstance(frame, core.Frame))
self.assertTrue(isinstance(frame[0], smb2.Smb2))
self.assertTrue(isinstance(frame[0][0], smb2.NegotiateResponse))
self.assertTrue(hasattr(frame, "buf"))
self.assertEqual(frame.len + 4, len(frame.buf))
callback_future.complete(True)
self.default_client.register_callback(model.EV_RES_POST_DESERIALIZE, cb)
conn = self.default_client.connect(self.server)
conn.negotiate()
self.assertTrue(callback_future.result(timeout=2))
