Source Code for Module pike.test.internal.test_callbacks

  1  import pike.core as core 
  2  import pike.model as model 
  3  import pike.smb2 as smb2 
  4  import pike.test 
  5   
  6 -class TestClientCallbacks(pike.test.PikeTest): 
  7 -    def test_pre_serialize(self): 
  8          callback_future = model.Future() 
  9          def cb(frame): 
 10              with callback_future: 
 11                  self.assertTrue(isinstance(frame, core.Frame)) 
 12                  self.assertTrue(isinstance(frame[0], smb2.Smb2)) 
 13                  self.assertTrue(isinstance(frame[0][0], smb2.NegotiateRequest)) 
 14                  self.assertFalse(hasattr(frame, "buf")) 
 15                  callback_future.complete(True) 
 16          self.default_client.register_callback(model.EV_REQ_PRE_SERIALIZE, cb) 
 17          conn = self.default_client.connect(self.server) 
 18          conn.negotiate() 
 19          self.assertTrue(callback_future.result(timeout=2)) 
 20   
 21 -    def test_post_serialize(self): 
 22          callback_future = model.Future() 
 23          def cb(frame): 
 24              with callback_future: 
 25                  self.assertTrue(isinstance(frame, core.Frame)) 
 26                  self.assertTrue(isinstance(frame[0], smb2.Smb2)) 
 27                  self.assertTrue(isinstance(frame[0][0], smb2.NegotiateRequest)) 
 28                  self.assertTrue(hasattr(frame, "buf")) 
 29                  self.assertEqual(frame.len + 4, len(frame.buf)) 
 30                  callback_future.complete(True) 
 31          self.default_client.register_callback(model.EV_REQ_POST_SERIALIZE, cb) 
 32          conn = self.default_client.connect(self.server) 
 33          conn.negotiate() 
 34          self.assertTrue(callback_future.result(timeout=2)) 
 35   
 36 -    def test_pre_post_send(self): 
 37          pre_callback_future = model.Future() 
 38          post_callback_future = model.Future() 
 39          expected_bytes = [] 
 40          def pre_cb(data): 
 41              with pre_callback_future: 
 42                  expected_bytes.insert(0, len(data)) 
 43                  self.assertGreater(expected_bytes[0], 16) 
 44                  pre_callback_future.complete(True) 
 45          def post_cb(bytes_written): 
 46              with post_callback_future: 
 47                  self.assertEqual(expected_bytes.pop(), bytes_written) 
 48                  if not expected_bytes: 
 49                      post_callback_future.complete(True) 
 50          self.default_client.register_callback(model.EV_REQ_PRE_SEND, pre_cb) 
 51          self.default_client.register_callback(model.EV_REQ_POST_SEND, post_cb) 
 52          conn = self.default_client.connect(self.server) 
 53          conn.negotiate() 
 54          self.assertTrue(pre_callback_future.result(timeout=2)) 
 55          self.assertTrue(post_callback_future.result(timeout=2)) 
 56   
 57 -    def test_pre_deserialize(self): 
 58          callback_future = model.Future() 
 59          def cb(data): 
 60              with callback_future: 
 61                  self.assertGreater(len(data), 4) 
 62                  callback_future.complete(True) 
 63          self.default_client.register_callback(model.EV_RES_PRE_DESERIALIZE, cb) 
 64          conn = self.default_client.connect(self.server) 
 65          conn.negotiate() 
 66          self.assertTrue(callback_future.result(timeout=2)) 
 67   
 68 -    def test_post_deserialize(self): 
 69          callback_future = model.Future() 
 70          def cb(frame): 
 71              with callback_future: 
 72                  self.assertTrue(isinstance(frame, core.Frame)) 
 73                  self.assertTrue(isinstance(frame[0], smb2.Smb2)) 
 74                  self.assertTrue(isinstance(frame[0][0], smb2.NegotiateResponse)) 
 75                  self.assertTrue(hasattr(frame, "buf")) 
 76                  self.assertEqual(frame.len + 4, len(frame.buf)) 
 77                  callback_future.complete(True) 
 78          self.default_client.register_callback(model.EV_RES_POST_DESERIALIZE, cb) 
 79          conn = self.default_client.connect(self.server) 
 80          conn.negotiate() 
 81          self.assertTrue(callback_future.result(timeout=2)) 
 82   
 83 -    def test_pre_post_recv(self): 
 84          pre_callback_future = model.Future() 
 85          post_callback_future = model.Future() 
 86          expected_bytes = [] 
 87          expected_rounds = [1,2] 
 88          def pre_cb(read_bytes): 
 89              with pre_callback_future: 
 90                  expected_bytes.insert(0, read_bytes) 
 91                  self.assertGreater(read_bytes, 3) 
 92                  pre_callback_future.complete(True) 
 93          def post_cb(data): 
 94              with post_callback_future: 
 95                  self.assertEqual(expected_bytes.pop(), len(data)) 
 96                  expected_rounds.pop() 
 97                  if not expected_rounds: 
 98                      post_callback_future.complete(True) 
 99          self.default_client.register_callback(model.EV_RES_PRE_RECV, pre_cb) 
100          self.default_client.register_callback(model.EV_RES_POST_RECV, post_cb) 
101          conn = self.default_client.connect(self.server) 
102          conn.negotiate() 
103          self.assertTrue(pre_callback_future.result(timeout=2)) 
104          self.assertTrue(post_callback_future.result(timeout=2)) 
105   
106 -class TestConnectionCallbacks(pike.test.PikeTest): 
107 -    def test_pre_serialize(self): 
108          callback_future = model.Future() 
109          def cb(frame): 
110              with callback_future: 
111                  self.assertTrue(isinstance(frame, core.Frame)) 
112                  self.assertTrue(isinstance(frame[0], smb2.Smb2)) 
113                  self.assertTrue(isinstance(frame[0][0], smb2.NegotiateRequest)) 
114                  self.assertFalse(hasattr(frame, "buf")) 
115                  callback_future.complete(True) 
116          conn = self.default_client.connect(self.server) 
117          conn.register_callback(model.EV_REQ_PRE_SERIALIZE, cb) 
118          conn.negotiate() 
119          self.assertTrue(callback_future.result(timeout=2)) 
120   
121 -    def test_post_serialize(self): 
122          callback_future = model.Future() 
123          def cb(frame): 
124              with callback_future: 
125                  self.assertTrue(isinstance(frame, core.Frame)) 
126                  self.assertTrue(isinstance(frame[0], smb2.Smb2)) 
127                  self.assertTrue(isinstance(frame[0][0], smb2.NegotiateRequest)) 
128                  self.assertTrue(hasattr(frame, "buf")) 
129                  self.assertEqual(frame.len + 4, len(frame.buf)) 
130                  callback_future.complete(True) 
131          conn = self.default_client.connect(self.server) 
132          conn.register_callback(model.EV_REQ_POST_SERIALIZE, cb) 
133          conn.negotiate() 
134          self.assertTrue(callback_future.result(timeout=2)) 
135   
136 -    def test_pre_post_send(self): 
137          pre_callback_future = model.Future() 
138          post_callback_future = model.Future() 
139          expected_bytes = [] 
140          def pre_cb(data): 
141              with pre_callback_future: 
142                  expected_bytes.insert(0, len(data)) 
143                  self.assertGreater(expected_bytes[0], 16) 
144                  pre_callback_future.complete(True) 
145          def post_cb(bytes_written): 
146              with post_callback_future: 
147                  self.assertEqual(expected_bytes.pop(), bytes_written) 
148                  if not expected_bytes: 
149                      post_callback_future.complete(True) 
150          conn = self.default_client.connect(self.server) 
151          conn.register_callback(model.EV_REQ_PRE_SEND, pre_cb) 
152          conn.register_callback(model.EV_REQ_POST_SEND, post_cb) 
153          conn.negotiate() 
154          self.assertTrue(pre_callback_future.result(timeout=2)) 
155          self.assertTrue(post_callback_future.result(timeout=2)) 
156   
157 -    def test_pre_deserialize(self): 
158          callback_future = model.Future() 
159          def cb(data): 
160              with callback_future: 
161                  self.assertGreater(len(data), 4) 
162                  callback_future.complete(True) 
163          conn = self.default_client.connect(self.server) 
164          conn.register_callback(model.EV_RES_PRE_DESERIALIZE, cb) 
165          conn.negotiate() 
166          self.assertTrue(callback_future.result(timeout=2)) 
167   
168 -    def test_post_deserialize(self): 
169          callback_future = model.Future() 
170          def cb(frame): 
171              with callback_future: 
172                  self.assertTrue(isinstance(frame, core.Frame)) 
173                  self.assertTrue(isinstance(frame[0], smb2.Smb2)) 
174                  self.assertTrue(isinstance(frame[0][0], smb2.NegotiateResponse)) 
175                  self.assertTrue(hasattr(frame, "buf")) 
176                  self.assertEqual(frame.len + 4, len(frame.buf)) 
177                  callback_future.complete(True) 
178          conn = self.default_client.connect(self.server) 
179          conn.register_callback(model.EV_RES_POST_DESERIALIZE, cb) 
180          conn.negotiate() 
181          self.assertTrue(callback_future.result(timeout=2)) 
182   
183 -    def test_pre_post_recv(self): 
184          pre_callback_future = model.Future() 
185          post_callback_future = model.Future() 
186          expected_bytes = [] 
187          expected_rounds = [1,2] 
188          def pre_cb(read_bytes): 
189              with pre_callback_future: 
190                  expected_bytes.insert(0, read_bytes) 
191                  self.assertGreater(read_bytes, 3) 
192                  pre_callback_future.complete(True) 
193          def post_cb(data): 
194              with post_callback_future: 
195                  self.assertEqual(expected_bytes.pop(), len(data)) 
196                  expected_rounds.pop() 
197                  if not expected_rounds: 
198                      post_callback_future.complete(True) 
199          conn = self.default_client.connect(self.server) 
200          conn.register_callback(model.EV_RES_PRE_RECV, pre_cb) 
201          conn.register_callback(model.EV_RES_POST_RECV, post_cb) 
202          conn.negotiate() 
203          self.assertTrue(pre_callback_future.result(timeout=2)) 
204          self.assertTrue(post_callback_future.result(timeout=2)) 
205