PEP8: fix E251: unexpected spaces around keyword / parameter equals
[amitay/samba.git] / python / samba / tests / dcerpc / raw_testcase.py
1 # Unix SMB/CIFS implementation.
2 # Copyright (C) Jelmer Vernooij <jelmer@samba.org> 2007-2010
3 # Copyright (C) Stefan Metzmacher 2014,2015
4 #
5 # This program is free software; you can redistribute it and/or modify
6 # it under the terms of the GNU General Public License as published by
7 # the Free Software Foundation; either version 3 of the License, or
8 # (at your option) any later version.
9 #
10 # This program is distributed in the hope that it will be useful,
11 # but WITHOUT ANY WARRANTY; without even the implied warranty of
12 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
13 # GNU General Public License for more details.
14 #
15 # You should have received a copy of the GNU General Public License
16 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
17 #
18
19 import sys
20 import socket
21 import struct
22 import samba.dcerpc.dcerpc
23 import samba.dcerpc.base
24 import samba.dcerpc.epmapper
25 import samba.tests
26 from samba import gensec
27 from samba.credentials import Credentials
28 from samba.tests import TestCase
29 from samba.ndr import ndr_pack, ndr_unpack, ndr_unpack_out
30
31 class RawDCERPCTest(TestCase):
32     """A raw DCE/RPC Test case."""
33
34     def _disconnect(self, reason):
35         if self.s is None:
36             return
37         self.s.close()
38         self.s = None
39         if self.do_hexdump:
40             sys.stderr.write("disconnect[%s]\n" % reason)
41
42     def connect(self):
43         try:
44             self.a = socket.getaddrinfo(self.host, self.tcp_port, socket.AF_UNSPEC,
45                                         socket.SOCK_STREAM, socket.SOL_TCP,
46                                         0)
47             self.s = socket.socket(self.a[0][0], self.a[0][1], self.a[0][2])
48             self.s.settimeout(10)
49             self.s.connect(self.a[0][4])
50         except socket.error as e:
51             self.s.close()
52             raise
53         except IOError as e:
54             self.s.close()
55             raise
56         except Exception as e:
57             raise
58         finally:
59             pass
60
61     def setUp(self):
62         super(RawDCERPCTest, self).setUp()
63         self.do_ndr_print = False
64         self.do_hexdump = False
65
66         self.host = samba.tests.env_get_var_value('SERVER')
67         self.target_hostname = samba.tests.env_get_var_value('TARGET_HOSTNAME', allow_missing=True)
68         if self.target_hostname is None:
69             self.target_hostname = self.host
70         self.tcp_port = 135
71
72         self.settings = {}
73         self.settings["lp_ctx"] = self.lp_ctx = samba.tests.env_loadparm()
74         self.settings["target_hostname"] = self.target_hostname
75
76         self.connect()
77
78     def noop(self):
79         return
80
81     def second_connection(self, tcp_port=None):
82         c = RawDCERPCTest(methodName='noop')
83         c.do_ndr_print = self.do_ndr_print
84         c.do_hexdump = self.do_hexdump
85
86         c.host = self.host
87         c.target_hostname = self.target_hostname
88         if tcp_port is not None:
89             c.tcp_port = tcp_port
90         else:
91             c.tcp_port = self.tcp_port
92
93         c.settings = self.settings
94
95         c.connect()
96         return c
97
98     def get_user_creds(self):
99         c = Credentials()
100         c.guess()
101         username = samba.tests.env_get_var_value('USERNAME')
102         password = samba.tests.env_get_var_value('PASSWORD')
103         c.set_username(username)
104         c.set_password(password)
105         return c
106
107     def get_anon_creds(self):
108         c = Credentials()
109         c.set_anonymous()
110         return c
111
112     def get_auth_context_creds(self, creds, auth_type, auth_level,
113                                auth_context_id,
114                                g_auth_level=None):
115
116         if g_auth_level is None:
117             g_auth_level = auth_level
118
119         g = gensec.Security.start_client(self.settings)
120         g.set_credentials(creds)
121         g.want_feature(gensec.FEATURE_DCE_STYLE)
122         g.start_mech_by_authtype(auth_type, g_auth_level)
123
124         auth_context = {}
125         auth_context["auth_type"] = auth_type
126         auth_context["auth_level"] = auth_level
127         auth_context["auth_context_id"] = auth_context_id
128         auth_context["g_auth_level"] = g_auth_level
129         auth_context["gensec"] = g
130
131         return auth_context
132
133     def do_generic_bind(self, ctx, auth_context=None,
134                         pfc_flags=samba.dcerpc.dcerpc.DCERPC_PFC_FLAG_FIRST |
135                         samba.dcerpc.dcerpc.DCERPC_PFC_FLAG_LAST,
136                         assoc_group_id=0, call_id=0,
137                         nak_reason=None, alter_fault=None):
138         ctx_list = [ctx]
139
140         if auth_context is not None:
141             from_server = ""
142             (finished, to_server) = auth_context["gensec"].update(from_server)
143             self.assertFalse(finished)
144
145             auth_info = self.generate_auth(auth_type=auth_context["auth_type"],
146                                            auth_level=auth_context["auth_level"],
147                                            auth_context_id=auth_context["auth_context_id"],
148                                            auth_blob=to_server)
149         else:
150             auth_info = ""
151
152         req = self.generate_bind(call_id=call_id,
153                                  pfc_flags=pfc_flags,
154                                  ctx_list=ctx_list,
155                                  assoc_group_id=assoc_group_id,
156                                  auth_info=auth_info)
157         self.send_pdu(req)
158         rep = self.recv_pdu()
159         if nak_reason is not None:
160             self.verify_pdu(rep, samba.dcerpc.dcerpc.DCERPC_PKT_BIND_NAK, req.call_id,
161                             auth_length=0)
162             self.assertEquals(rep.u.reject_reason, nak_reason)
163             self.assertEquals(rep.u.num_versions, 1)
164             self.assertEquals(rep.u.versions[0].rpc_vers, req.rpc_vers)
165             self.assertEquals(rep.u.versions[0].rpc_vers_minor, req.rpc_vers_minor)
166             self.assertEquals(len(rep.u._pad), 3)
167             self.assertEquals(rep.u._pad, '\0' * 3)
168             return
169         self.verify_pdu(rep, samba.dcerpc.dcerpc.DCERPC_PKT_BIND_ACK, req.call_id,
170                         pfc_flags=pfc_flags)
171         self.assertEquals(rep.u.max_xmit_frag, req.u.max_xmit_frag)
172         self.assertEquals(rep.u.max_recv_frag, req.u.max_recv_frag)
173         if assoc_group_id != 0:
174             self.assertEquals(rep.u.assoc_group_id, assoc_group_id)
175         else:
176             self.assertNotEquals(rep.u.assoc_group_id, 0)
177             assoc_group_id = rep.u.assoc_group_id
178         port_str = "%d" % self.tcp_port
179         port_len = len(port_str) + 1
180         mod_len = (2 + port_len) % 4
181         if mod_len != 0:
182             port_pad = 4 - mod_len
183         else:
184             port_pad = 0
185         self.assertEquals(rep.u.secondary_address_size, port_len)
186         self.assertEquals(rep.u.secondary_address, port_str)
187         self.assertEquals(len(rep.u._pad1), port_pad)
188         # sometimes windows sends random bytes
189         # self.assertEquals(rep.u._pad1, '\0' * port_pad)
190         self.assertEquals(rep.u.num_results, 1)
191         self.assertEquals(rep.u.ctx_list[0].result,
192                           samba.dcerpc.dcerpc.DCERPC_BIND_ACK_RESULT_ACCEPTANCE)
193         self.assertEquals(rep.u.ctx_list[0].reason,
194                           samba.dcerpc.dcerpc.DCERPC_BIND_ACK_REASON_NOT_SPECIFIED)
195         self.assertNDRSyntaxEquals(rep.u.ctx_list[0].syntax, ctx.transfer_syntaxes[0])
196         ack = rep
197         if auth_context is None:
198             self.assertEquals(rep.auth_length, 0)
199             self.assertEquals(len(rep.u.auth_info), 0)
200             return ack
201         self.assertNotEquals(rep.auth_length, 0)
202         self.assertGreater(len(rep.u.auth_info), samba.dcerpc.dcerpc.DCERPC_AUTH_TRAILER_LENGTH)
203         self.assertEquals(rep.auth_length, len(rep.u.auth_info) - samba.dcerpc.dcerpc.DCERPC_AUTH_TRAILER_LENGTH)
204
205         a = self.parse_auth(rep.u.auth_info)
206
207         from_server = a.credentials
208         (finished, to_server) = auth_context["gensec"].update(from_server)
209         self.assertFalse(finished)
210
211         auth_info = self.generate_auth(auth_type=auth_context["auth_type"],
212                                        auth_level=auth_context["auth_level"],
213                                        auth_context_id=auth_context["auth_context_id"],
214                                        auth_blob=to_server)
215         req = self.generate_alter(call_id=call_id,
216                                   ctx_list=ctx_list,
217                                   assoc_group_id=0xffffffff -assoc_group_id,
218                                   auth_info=auth_info)
219         self.send_pdu(req)
220         rep = self.recv_pdu()
221         if alter_fault is not None:
222             self.verify_pdu(rep, samba.dcerpc.dcerpc.DCERPC_PKT_FAULT, req.call_id,
223                             pfc_flags=req.pfc_flags |
224                             samba.dcerpc.dcerpc.DCERPC_PFC_FLAG_DID_NOT_EXECUTE,
225                             auth_length=0)
226             self.assertNotEquals(rep.u.alloc_hint, 0)
227             self.assertEquals(rep.u.context_id, 0)
228             self.assertEquals(rep.u.cancel_count, 0)
229             self.assertEquals(rep.u.flags, 0)
230             self.assertEquals(rep.u.status, alter_fault)
231             self.assertEquals(rep.u.reserved, 0)
232             self.assertEquals(len(rep.u.error_and_verifier), 0)
233             return None
234         self.verify_pdu(rep, samba.dcerpc.dcerpc.DCERPC_PKT_ALTER_RESP, req.call_id)
235         self.assertEquals(rep.u.max_xmit_frag, req.u.max_xmit_frag)
236         self.assertEquals(rep.u.max_recv_frag, req.u.max_recv_frag)
237         self.assertEquals(rep.u.assoc_group_id, assoc_group_id)
238         self.assertEquals(rep.u.secondary_address_size, 0)
239         self.assertEquals(rep.u.secondary_address, '')
240         self.assertEquals(len(rep.u._pad1), 2)
241         # sometimes windows sends random bytes
242         # self.assertEquals(rep.u._pad1, '\0' * 2)
243         self.assertEquals(rep.u.num_results, 1)
244         self.assertEquals(rep.u.ctx_list[0].result,
245                           samba.dcerpc.dcerpc.DCERPC_BIND_ACK_RESULT_ACCEPTANCE)
246         self.assertEquals(rep.u.ctx_list[0].reason,
247                           samba.dcerpc.dcerpc.DCERPC_BIND_ACK_REASON_NOT_SPECIFIED)
248         self.assertNDRSyntaxEquals(rep.u.ctx_list[0].syntax, ctx.transfer_syntaxes[0])
249         self.assertNotEquals(rep.auth_length, 0)
250         self.assertGreater(len(rep.u.auth_info), samba.dcerpc.dcerpc.DCERPC_AUTH_TRAILER_LENGTH)
251         self.assertEquals(rep.auth_length, len(rep.u.auth_info) - samba.dcerpc.dcerpc.DCERPC_AUTH_TRAILER_LENGTH)
252
253         a = self.parse_auth(rep.u.auth_info)
254
255         from_server = a.credentials
256         (finished, to_server) = auth_context["gensec"].update(from_server)
257         self.assertTrue(finished)
258
259         return ack
260
261     def prepare_presentation(self, abstract, transfer, object=None,
262                              context_id=0xffff, epmap=False, auth_context=None,
263                              pfc_flags=samba.dcerpc.dcerpc.DCERPC_PFC_FLAG_FIRST |
264                              samba.dcerpc.dcerpc.DCERPC_PFC_FLAG_LAST,
265                              assoc_group_id=0,
266                              return_ack=False):
267         if epmap:
268             self.epmap_reconnect(abstract, transfer=transfer, object=object)
269
270         tsf1_list = [transfer]
271         ctx = samba.dcerpc.dcerpc.ctx_list()
272         ctx.context_id = context_id
273         ctx.num_transfer_syntaxes = len(tsf1_list)
274         ctx.abstract_syntax = abstract
275         ctx.transfer_syntaxes = tsf1_list
276
277         ack = self.do_generic_bind(ctx=ctx,
278                                    auth_context=auth_context,
279                                    pfc_flags=pfc_flags,
280                                    assoc_group_id=assoc_group_id)
281         if ack is None:
282             ctx = None
283
284         if return_ack:
285             return (ctx, ack)
286         return ctx
287
288     def do_single_request(self, call_id, ctx, io,
289                           auth_context=None,
290                           object=None,
291                           bigendian=False, ndr64=False,
292                           allow_remaining=False,
293                           send_req=True,
294                           recv_rep=True,
295                           fault_pfc_flags=samba.dcerpc.dcerpc.DCERPC_PFC_FLAG_FIRST |
296                           samba.dcerpc.dcerpc.DCERPC_PFC_FLAG_LAST,
297                           fault_status=None,
298                           fault_context_id=None,
299                           timeout=None,
300                           ndr_print=None,
301                           hexdump=None):
302
303         if fault_context_id is None:
304             fault_context_id = ctx.context_id
305
306         if ndr_print is None:
307             ndr_print = self.do_ndr_print
308         if hexdump is None:
309             hexdump = self.do_hexdump
310
311         if send_req:
312             if ndr_print:
313                 sys.stderr.write("in: %s" % samba.ndr.ndr_print_in(io))
314             stub_in = samba.ndr.ndr_pack_in(io, bigendian=bigendian, ndr64=ndr64)
315             if hexdump:
316                 sys.stderr.write("stub_in: %d\n%s" % (len(stub_in), self.hexdump(stub_in)))
317         else:
318             # only used for sig_size calculation
319             stub_in = '\xff' * samba.dcerpc.dcerpc.DCERPC_AUTH_PAD_ALIGNMENT
320
321         sig_size = 0
322         if auth_context is not None:
323             mod_len = len(stub_in) % samba.dcerpc.dcerpc.DCERPC_AUTH_PAD_ALIGNMENT
324             auth_pad_length = 0
325             if mod_len > 0:
326                 auth_pad_length = samba.dcerpc.dcerpc.DCERPC_AUTH_PAD_ALIGNMENT - mod_len
327             stub_in += '\x00' * auth_pad_length
328
329             if auth_context["g_auth_level"] >= samba.dcerpc.dcerpc.DCERPC_AUTH_LEVEL_PACKET:
330                 sig_size = auth_context["gensec"].sig_size(len(stub_in))
331             else:
332                 sig_size = 16
333
334             zero_sig = "\x00" * sig_size
335             auth_info = self.generate_auth(auth_type=auth_context["auth_type"],
336                                            auth_level=auth_context["auth_level"],
337                                            auth_pad_length=auth_pad_length,
338                                            auth_context_id=auth_context["auth_context_id"],
339                                            auth_blob=zero_sig)
340         else:
341             auth_info = ""
342
343         pfc_flags = samba.dcerpc.dcerpc.DCERPC_PFC_FLAG_FIRST
344         pfc_flags |= samba.dcerpc.dcerpc.DCERPC_PFC_FLAG_LAST
345         if object is not None:
346             pfc_flags |= samba.dcerpc.dcerpc.DCERPC_PFC_FLAG_OBJECT_UUID
347
348         req = self.generate_request(call_id=call_id,
349                                     context_id=ctx.context_id,
350                                     pfc_flags=pfc_flags,
351                                     object=object,
352                                     opnum=io.opnum(),
353                                     stub=stub_in,
354                                     auth_info=auth_info)
355
356         if send_req:
357             if sig_size != 0 and auth_context["auth_level"] >= samba.dcerpc.dcerpc.DCERPC_AUTH_LEVEL_PACKET:
358                 req_blob = samba.ndr.ndr_pack(req)
359                 ofs_stub = samba.dcerpc.dcerpc.DCERPC_REQUEST_LENGTH
360                 ofs_sig = len(req_blob) - req.auth_length
361                 ofs_trailer = ofs_sig - samba.dcerpc.dcerpc.DCERPC_AUTH_TRAILER_LENGTH
362                 req_data = req_blob[ofs_stub:ofs_trailer]
363                 req_whole = req_blob[0:ofs_sig]
364                 sig = auth_context["gensec"].sign_packet(req_data, req_whole)
365                 auth_info = self.generate_auth(auth_type=auth_context["auth_type"],
366                                                auth_level=auth_context["auth_level"],
367                                                auth_pad_length=auth_pad_length,
368                                                auth_context_id=auth_context["auth_context_id"],
369                                                auth_blob=sig)
370                 req = self.generate_request(call_id=call_id,
371                                             context_id=ctx.context_id,
372                                             pfc_flags=pfc_flags,
373                                             object=object,
374                                             opnum=io.opnum(),
375                                             stub=stub_in,
376                                             auth_info=auth_info)
377             self.send_pdu(req, ndr_print=ndr_print, hexdump=hexdump)
378         if recv_rep:
379             (rep, rep_blob) = self.recv_pdu_raw(timeout=timeout,
380                                                 ndr_print=ndr_print,
381                                                 hexdump=hexdump)
382             if fault_status:
383                 self.verify_pdu(rep, samba.dcerpc.dcerpc.DCERPC_PKT_FAULT, req.call_id,
384                                 pfc_flags=fault_pfc_flags, auth_length=0)
385                 self.assertNotEquals(rep.u.alloc_hint, 0)
386                 self.assertEquals(rep.u.context_id, fault_context_id)
387                 self.assertEquals(rep.u.cancel_count, 0)
388                 self.assertEquals(rep.u.flags, 0)
389                 self.assertEquals(rep.u.status, fault_status)
390                 self.assertEquals(rep.u.reserved, 0)
391                 self.assertEquals(len(rep.u.error_and_verifier), 0)
392                 return
393
394             self.verify_pdu(rep, samba.dcerpc.dcerpc.DCERPC_PKT_RESPONSE, req.call_id,
395                             auth_length=sig_size)
396             self.assertNotEquals(rep.u.alloc_hint, 0)
397             self.assertEquals(rep.u.context_id, req.u.context_id & 0xff)
398             self.assertEquals(rep.u.cancel_count, 0)
399             self.assertGreaterEqual(len(rep.u.stub_and_verifier), rep.u.alloc_hint)
400             if sig_size != 0:
401
402                 ofs_stub = samba.dcerpc.dcerpc.DCERPC_REQUEST_LENGTH
403                 ofs_sig = rep.frag_length - rep.auth_length
404                 ofs_trailer = ofs_sig - samba.dcerpc.dcerpc.DCERPC_AUTH_TRAILER_LENGTH
405                 rep_data = rep_blob[ofs_stub:ofs_trailer]
406                 rep_whole = rep_blob[0:ofs_sig]
407                 rep_sig = rep_blob[ofs_sig:]
408                 rep_auth_info_blob = rep_blob[ofs_trailer:]
409
410                 rep_auth_info = self.parse_auth(rep_auth_info_blob)
411                 self.assertEquals(rep_auth_info.auth_type, auth_context["auth_type"])
412                 self.assertEquals(rep_auth_info.auth_level, auth_context["auth_level"])
413                 self.assertLessEqual(rep_auth_info.auth_pad_length, len(rep_data))
414                 self.assertEquals(rep_auth_info.auth_reserved, 0)
415                 self.assertEquals(rep_auth_info.auth_context_id, auth_context["auth_context_id"])
416                 self.assertEquals(rep_auth_info.credentials, rep_sig)
417
418                 if auth_context["auth_level"] >= samba.dcerpc.dcerpc.DCERPC_AUTH_LEVEL_PACKET:
419                     auth_context["gensec"].check_packet(rep_data, rep_whole, rep_sig)
420
421                 stub_out = rep_data[0:-rep_auth_info.auth_pad_length]
422             else:
423                 stub_out = rep.u.stub_and_verifier
424
425             if hexdump:
426                 sys.stderr.write("stub_out: %d\n%s" % (len(stub_out), self.hexdump(stub_out)))
427             ndr_unpack_out(io, stub_out, bigendian=bigendian, ndr64=ndr64,
428                            allow_remaining=allow_remaining)
429             if ndr_print:
430                 sys.stderr.write("out: %s" % samba.ndr.ndr_print_out(io))
431
432     def epmap_reconnect(self, abstract, transfer=None, object=None):
433         ndr32 = samba.dcerpc.base.transfer_syntax_ndr()
434
435         if transfer is None:
436             transfer = ndr32
437
438         if object is None:
439             object = samba.dcerpc.misc.GUID()
440
441         ctx = self.prepare_presentation(samba.dcerpc.epmapper.abstract_syntax(),
442                                         transfer, context_id=0)
443
444         data1 = ndr_pack(abstract)
445         lhs1 = samba.dcerpc.epmapper.epm_lhs()
446         lhs1.protocol = samba.dcerpc.epmapper.EPM_PROTOCOL_UUID
447         lhs1.lhs_data = data1[:18]
448         rhs1 = samba.dcerpc.epmapper.epm_rhs_uuid()
449         rhs1.unknown = data1[18:]
450         floor1 = samba.dcerpc.epmapper.epm_floor()
451         floor1.lhs = lhs1
452         floor1.rhs = rhs1
453         data2 = ndr_pack(transfer)
454         lhs2 = samba.dcerpc.epmapper.epm_lhs()
455         lhs2.protocol = samba.dcerpc.epmapper.EPM_PROTOCOL_UUID
456         lhs2.lhs_data = data2[:18]
457         rhs2 = samba.dcerpc.epmapper.epm_rhs_uuid()
458         rhs2.unknown = data1[18:]
459         floor2 = samba.dcerpc.epmapper.epm_floor()
460         floor2.lhs = lhs2
461         floor2.rhs = rhs2
462         lhs3 = samba.dcerpc.epmapper.epm_lhs()
463         lhs3.protocol = samba.dcerpc.epmapper.EPM_PROTOCOL_NCACN
464         lhs3.lhs_data = ""
465         floor3 = samba.dcerpc.epmapper.epm_floor()
466         floor3.lhs = lhs3
467         floor3.rhs.minor_version = 0
468         lhs4 = samba.dcerpc.epmapper.epm_lhs()
469         lhs4.protocol = samba.dcerpc.epmapper.EPM_PROTOCOL_TCP
470         lhs4.lhs_data = ""
471         floor4 = samba.dcerpc.epmapper.epm_floor()
472         floor4.lhs = lhs4
473         floor4.rhs.port = self.tcp_port
474         lhs5 = samba.dcerpc.epmapper.epm_lhs()
475         lhs5.protocol = samba.dcerpc.epmapper.EPM_PROTOCOL_IP
476         lhs5.lhs_data = ""
477         floor5 = samba.dcerpc.epmapper.epm_floor()
478         floor5.lhs = lhs5
479         floor5.rhs.ipaddr = "0.0.0.0"
480
481         floors = [floor1, floor2, floor3, floor4, floor5]
482         req_tower = samba.dcerpc.epmapper.epm_tower()
483         req_tower.num_floors = len(floors)
484         req_tower.floors = floors
485         req_twr = samba.dcerpc.epmapper.epm_twr_t()
486         req_twr.tower = req_tower
487
488         epm_map = samba.dcerpc.epmapper.epm_Map()
489         epm_map.in_object = object
490         epm_map.in_map_tower = req_twr
491         epm_map.in_entry_handle = samba.dcerpc.misc.policy_handle()
492         epm_map.in_max_towers = 4
493
494         self.do_single_request(call_id=2, ctx=ctx, io=epm_map)
495
496         self.assertGreaterEqual(epm_map.out_num_towers, 1)
497         rep_twr = epm_map.out_towers[0].twr
498         self.assertIsNotNone(rep_twr)
499         self.assertEqual(rep_twr.tower_length, 75)
500         self.assertEqual(rep_twr.tower.num_floors, 5)
501         self.assertEqual(len(rep_twr.tower.floors), 5)
502         self.assertEqual(rep_twr.tower.floors[3].lhs.protocol,
503                          samba.dcerpc.epmapper.EPM_PROTOCOL_TCP)
504         self.assertEqual(rep_twr.tower.floors[3].lhs.protocol,
505                          samba.dcerpc.epmapper.EPM_PROTOCOL_TCP)
506
507         # reconnect to the given port
508         self._disconnect("epmap_reconnect")
509         self.tcp_port = rep_twr.tower.floors[3].rhs.port
510         self.connect()
511
512     def send_pdu(self, req, ndr_print=None, hexdump=None):
513         if ndr_print is None:
514             ndr_print = self.do_ndr_print
515         if hexdump is None:
516             hexdump = self.do_hexdump
517         try:
518             req_pdu = ndr_pack(req)
519             if ndr_print:
520                 sys.stderr.write("send_pdu: %s" % samba.ndr.ndr_print(req))
521             if hexdump:
522                 sys.stderr.write("send_pdu: %d\n%s" % (len(req_pdu), self.hexdump(req_pdu)))
523             while True:
524                 sent = self.s.send(req_pdu, 0)
525                 if sent == len(req_pdu):
526                     break
527                 req_pdu = req_pdu[sent:]
528         except socket.error as e:
529             self._disconnect("send_pdu: %s" % e)
530             raise
531         except IOError as e:
532             self._disconnect("send_pdu: %s" % e)
533             raise
534         finally:
535             pass
536
537     def recv_raw(self, hexdump=None, timeout=None):
538         rep_pdu = None
539         if hexdump is None:
540             hexdump = self.do_hexdump
541         try:
542             if timeout is not None:
543                 self.s.settimeout(timeout)
544             rep_pdu = self.s.recv(0xffff, 0)
545             self.s.settimeout(10)
546             if len(rep_pdu) == 0:
547                 self._disconnect("recv_raw: EOF")
548                 return None
549             if hexdump:
550                 sys.stderr.write("recv_raw: %d\n%s" % (len(rep_pdu), self.hexdump(rep_pdu)))
551         except socket.timeout as e:
552             self.s.settimeout(10)
553             sys.stderr.write("recv_raw: TIMEOUT\n")
554             pass
555         except socket.error as e:
556             self._disconnect("recv_raw: %s" % e)
557             raise
558         except IOError as e:
559             self._disconnect("recv_raw: %s" % e)
560             raise
561         finally:
562             pass
563         return rep_pdu
564
565     def recv_pdu_raw(self, ndr_print=None, hexdump=None, timeout=None):
566         rep_pdu = None
567         rep = None
568         if ndr_print is None:
569             ndr_print = self.do_ndr_print
570         if hexdump is None:
571             hexdump = self.do_hexdump
572         try:
573             rep_pdu = self.recv_raw(hexdump=hexdump, timeout=timeout)
574             if rep_pdu is None:
575                 return (None, None)
576             rep = ndr_unpack(samba.dcerpc.dcerpc.ncacn_packet, rep_pdu, allow_remaining=True)
577             if ndr_print:
578                 sys.stderr.write("recv_pdu: %s" % samba.ndr.ndr_print(rep))
579             self.assertEqual(rep.frag_length, len(rep_pdu))
580         finally:
581             pass
582         return (rep, rep_pdu)
583
584     def recv_pdu(self, ndr_print=None, hexdump=None, timeout=None):
585         (rep, rep_pdu) = self.recv_pdu_raw(ndr_print=ndr_print,
586                                            hexdump=hexdump,
587                                            timeout=timeout)
588         return rep
589
590     def generate_auth(self,
591                       auth_type=None,
592                       auth_level=None,
593                       auth_pad_length=0,
594                       auth_context_id=None,
595                       auth_blob=None,
596                       ndr_print=None, hexdump=None):
597         if ndr_print is None:
598             ndr_print = self.do_ndr_print
599         if hexdump is None:
600             hexdump = self.do_hexdump
601
602         if auth_type is not None:
603             a = samba.dcerpc.dcerpc.auth()
604             a.auth_type = auth_type
605             a.auth_level = auth_level
606             a.auth_pad_length = auth_pad_length
607             a.auth_context_id = auth_context_id
608             a.credentials = auth_blob
609
610             ai = ndr_pack(a)
611             if ndr_print:
612                 sys.stderr.write("generate_auth: %s" % samba.ndr.ndr_print(a))
613             if hexdump:
614                 sys.stderr.write("generate_auth: %d\n%s" % (len(ai), self.hexdump(ai)))
615         else:
616             ai = ""
617
618         return ai
619
620     def parse_auth(self, auth_info, ndr_print=None, hexdump=None):
621         if ndr_print is None:
622             ndr_print = self.do_ndr_print
623         if hexdump is None:
624             hexdump = self.do_hexdump
625
626         if (len(auth_info) <= samba.dcerpc.dcerpc.DCERPC_AUTH_TRAILER_LENGTH):
627             return None
628
629         if hexdump:
630             sys.stderr.write("parse_auth: %d\n%s" % (len(auth_info), self.hexdump(auth_info)))
631         a = ndr_unpack(samba.dcerpc.dcerpc.auth, auth_info, allow_remaining=True)
632         if ndr_print:
633             sys.stderr.write("parse_auth: %s" % samba.ndr.ndr_print(a))
634
635         return a
636
637     def generate_pdu(self, ptype, call_id, payload,
638                      rpc_vers=5,
639                      rpc_vers_minor=0,
640                      pfc_flags=samba.dcerpc.dcerpc.DCERPC_PFC_FLAG_FIRST |
641                      samba.dcerpc.dcerpc.DCERPC_PFC_FLAG_LAST,
642                      drep=[samba.dcerpc.dcerpc.DCERPC_DREP_LE, 0, 0, 0],
643                      ndr_print=None, hexdump=None):
644
645         if getattr(payload, 'auth_info', None):
646             ai = payload.auth_info
647         else:
648             ai = ""
649
650         p = samba.dcerpc.dcerpc.ncacn_packet()
651         p.rpc_vers = rpc_vers
652         p.rpc_vers_minor = rpc_vers_minor
653         p.ptype = ptype
654         p.pfc_flags = pfc_flags
655         p.drep = drep
656         p.frag_length = 0
657         if len(ai) > samba.dcerpc.dcerpc.DCERPC_AUTH_TRAILER_LENGTH:
658             p.auth_length = len(ai) - samba.dcerpc.dcerpc.DCERPC_AUTH_TRAILER_LENGTH
659         else:
660             p.auth_length = 0
661         p.call_id = call_id
662         p.u = payload
663
664         pdu = ndr_pack(p)
665         p.frag_length = len(pdu)
666
667         return p
668
669     def verify_pdu(self, p, ptype, call_id,
670                    rpc_vers=5,
671                    rpc_vers_minor=0,
672                    pfc_flags=samba.dcerpc.dcerpc.DCERPC_PFC_FLAG_FIRST |
673                    samba.dcerpc.dcerpc.DCERPC_PFC_FLAG_LAST,
674                    drep=[samba.dcerpc.dcerpc.DCERPC_DREP_LE, 0, 0, 0],
675                    auth_length=None):
676
677         self.assertIsNotNone(p, "No valid pdu")
678
679         if getattr(p.u, 'auth_info', None):
680             ai = p.u.auth_info
681         else:
682             ai = ""
683
684         self.assertEqual(p.rpc_vers, rpc_vers)
685         self.assertEqual(p.rpc_vers_minor, rpc_vers_minor)
686         self.assertEqual(p.ptype, ptype)
687         self.assertEqual(p.pfc_flags, pfc_flags)
688         self.assertEqual(p.drep, drep)
689         self.assertGreaterEqual(p.frag_length,
690                                 samba.dcerpc.dcerpc.DCERPC_NCACN_PAYLOAD_OFFSET)
691         if len(ai) > samba.dcerpc.dcerpc.DCERPC_AUTH_TRAILER_LENGTH:
692             self.assertEqual(p.auth_length,
693                              len(ai) - samba.dcerpc.dcerpc.DCERPC_AUTH_TRAILER_LENGTH)
694         elif auth_length is not None:
695             self.assertEqual(p.auth_length, auth_length)
696         else:
697             self.assertEqual(p.auth_length, 0)
698         self.assertEqual(p.call_id, call_id)
699
700         return
701
702     def generate_bind(self, call_id,
703                       pfc_flags=samba.dcerpc.dcerpc.DCERPC_PFC_FLAG_FIRST |
704                       samba.dcerpc.dcerpc.DCERPC_PFC_FLAG_LAST,
705                       max_xmit_frag=5840,
706                       max_recv_frag=5840,
707                       assoc_group_id=0,
708                       ctx_list=[],
709                       auth_info="",
710                       ndr_print=None, hexdump=None):
711
712         b = samba.dcerpc.dcerpc.bind()
713         b.max_xmit_frag = max_xmit_frag
714         b.max_recv_frag = max_recv_frag
715         b.assoc_group_id = assoc_group_id
716         b.num_contexts = len(ctx_list)
717         b.ctx_list = ctx_list
718         b.auth_info = auth_info
719
720         p = self.generate_pdu(ptype=samba.dcerpc.dcerpc.DCERPC_PKT_BIND,
721                               pfc_flags=pfc_flags,
722                               call_id=call_id,
723                               payload=b,
724                               ndr_print=ndr_print, hexdump=hexdump)
725
726         return p
727
728     def generate_alter(self, call_id,
729                        pfc_flags=samba.dcerpc.dcerpc.DCERPC_PFC_FLAG_FIRST |
730                        samba.dcerpc.dcerpc.DCERPC_PFC_FLAG_LAST,
731                        max_xmit_frag=5840,
732                        max_recv_frag=5840,
733                        assoc_group_id=0,
734                        ctx_list=[],
735                        auth_info="",
736                        ndr_print=None, hexdump=None):
737
738         a = samba.dcerpc.dcerpc.bind()
739         a.max_xmit_frag = max_xmit_frag
740         a.max_recv_frag = max_recv_frag
741         a.assoc_group_id = assoc_group_id
742         a.num_contexts = len(ctx_list)
743         a.ctx_list = ctx_list
744         a.auth_info = auth_info
745
746         p = self.generate_pdu(ptype=samba.dcerpc.dcerpc.DCERPC_PKT_ALTER,
747                               pfc_flags=pfc_flags,
748                               call_id=call_id,
749                               payload=a,
750                               ndr_print=ndr_print, hexdump=hexdump)
751
752         return p
753
754     def generate_auth3(self, call_id,
755                        pfc_flags=samba.dcerpc.dcerpc.DCERPC_PFC_FLAG_FIRST |
756                        samba.dcerpc.dcerpc.DCERPC_PFC_FLAG_LAST,
757                        auth_info="",
758                        ndr_print=None, hexdump=None):
759
760         a = samba.dcerpc.dcerpc.auth3()
761         a.auth_info = auth_info
762
763         p = self.generate_pdu(ptype=samba.dcerpc.dcerpc.DCERPC_PKT_AUTH3,
764                               pfc_flags=pfc_flags,
765                               call_id=call_id,
766                               payload=a,
767                               ndr_print=ndr_print, hexdump=hexdump)
768
769         return p
770
771     def generate_request(self, call_id,
772                          pfc_flags=samba.dcerpc.dcerpc.DCERPC_PFC_FLAG_FIRST |
773                          samba.dcerpc.dcerpc.DCERPC_PFC_FLAG_LAST,
774                          alloc_hint=None,
775                          context_id=None,
776                          opnum=None,
777                          object=None,
778                          stub=None,
779                          auth_info="",
780                          ndr_print=None, hexdump=None):
781
782         if alloc_hint is None:
783             alloc_hint = len(stub)
784
785         r = samba.dcerpc.dcerpc.request()
786         r.alloc_hint = alloc_hint
787         r.context_id = context_id
788         r.opnum = opnum
789         if object is not None:
790             r.object = object
791         r.stub_and_verifier = stub + auth_info
792
793         p = self.generate_pdu(ptype=samba.dcerpc.dcerpc.DCERPC_PKT_REQUEST,
794                               pfc_flags=pfc_flags,
795                               call_id=call_id,
796                               payload=r,
797                               ndr_print=ndr_print, hexdump=hexdump)
798
799         if len(auth_info) > samba.dcerpc.dcerpc.DCERPC_AUTH_TRAILER_LENGTH:
800             p.auth_length = len(auth_info) - samba.dcerpc.dcerpc.DCERPC_AUTH_TRAILER_LENGTH
801
802         return p
803
804     def generate_co_cancel(self, call_id,
805                            pfc_flags=samba.dcerpc.dcerpc.DCERPC_PFC_FLAG_FIRST |
806                            samba.dcerpc.dcerpc.DCERPC_PFC_FLAG_LAST,
807                            auth_info="",
808                            ndr_print=None, hexdump=None):
809
810         c = samba.dcerpc.dcerpc.co_cancel()
811         c.auth_info = auth_info
812
813         p = self.generate_pdu(ptype=samba.dcerpc.dcerpc.DCERPC_PKT_CO_CANCEL,
814                               pfc_flags=pfc_flags,
815                               call_id=call_id,
816                               payload=c,
817                               ndr_print=ndr_print, hexdump=hexdump)
818
819         return p
820
821     def generate_orphaned(self, call_id,
822                           pfc_flags=samba.dcerpc.dcerpc.DCERPC_PFC_FLAG_FIRST |
823                           samba.dcerpc.dcerpc.DCERPC_PFC_FLAG_LAST,
824                           auth_info="",
825                           ndr_print=None, hexdump=None):
826
827         o = samba.dcerpc.dcerpc.orphaned()
828         o.auth_info = auth_info
829
830         p = self.generate_pdu(ptype=samba.dcerpc.dcerpc.DCERPC_PKT_ORPHANED,
831                               pfc_flags=pfc_flags,
832                               call_id=call_id,
833                               payload=o,
834                               ndr_print=ndr_print, hexdump=hexdump)
835
836         return p
837
838     def generate_shutdown(self, call_id,
839                           pfc_flags=samba.dcerpc.dcerpc.DCERPC_PFC_FLAG_FIRST |
840                           samba.dcerpc.dcerpc.DCERPC_PFC_FLAG_LAST,
841                           ndr_print=None, hexdump=None):
842
843         s = samba.dcerpc.dcerpc.shutdown()
844
845         p = self.generate_pdu(ptype=samba.dcerpc.dcerpc.DCERPC_PKT_SHUTDOWN,
846                               pfc_flags=pfc_flags,
847                               call_id=call_id,
848                               payload=s,
849                               ndr_print=ndr_print, hexdump=hexdump)
850
851         return p
852
853     def assertIsConnected(self):
854         self.assertIsNotNone(self.s, msg="Not connected")
855         return
856
857     def assertNotConnected(self):
858         self.assertIsNone(self.s, msg="Is connected")
859         return
860
861     def assertNDRSyntaxEquals(self, s1, s2):
862         self.assertEqual(s1.uuid, s2.uuid)
863         self.assertEqual(s1.if_version, s2.if_version)
864         return