pike.model

1 # 2 # Copyright (c) 2013, EMC Corporation 3 # All rights reserved. 4 # 5 # Redistribution and use in source and binary forms, with or without 6 # modification, are permitted provided that the following conditions are met: 7 # 8 # 1. Redistributions of source code must retain the above copyright notice, 9 # this list of conditions and the following disclaimer. 10 # 2. Redistributions in binary form must reproduce the above copyright notice, 11 # this list of conditions and the following disclaimer in the documentation 12 # and/or other materials provided with the distribution. 13 # 14 # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 15 # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 16 # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE 17 # ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE 18 # LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 19 # CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF 20 # SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS 21 # INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN 22 # CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) 23 # ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE 24 # POSSIBILITY OF SUCH DAMAGE. 25 # 26 # Module Name: 27 # 28 # model.py 29 # 30 # Abstract: 31 # 32 # Transport and object model 33 # 34 # Authors: Brian Koropoff (brian.koropoff@emc.com) 35 # 36 37 """ 38 SMB2 Object Model. 39 40 This module contains an implementation of the SMB2 client object model, 41 allowing channels, sessions, tree connections, opens, and leases 42 to be established and tracked. It provides convenience functions 43 for exercising common elements of the protocol without manually 44 constructing packets. 45 """ 46 47 import sys 48 import socket 49 import array 50 import struct 51 import random 52 import logging 53 import time 54 import operator 55 import contextlib 56 57 import auth 58 import core 59 import crypto 60 import netbios 61 import nttime 62 import smb2 63 import transport 64 import ntstatus 65 import digest 66 67 default_credit_request = 10 68 default_timeout = 30 69 trace = False

70 71 -def loop(timeout=None, count=None):

72 """ 73 wrapper for blocking on the underlying event loop for the given timeout 74 or given count of iterations 75 """ 76 if timeout is None: 77 timeout = default_timeout 78 transport.loop(timeout=timeout, count=count)

79

80 -class TimeoutError(Exception):

81 pass

82

83 -class StateError(Exception):

84 pass

85

86 -class CreditError(Exception):

87 pass

88

89 -class ResponseError(Exception):

90 - def __init__(self, response):

91 Exception.__init__(self, response.command, response.status) 92 self.response = response

93

94 -class Events(core.ValueEnum):

95 """ Events used for callback functions """ 96 EV_REQ_PRE_SERIALIZE = 0x1 # cb expects Netbios frame 97 EV_REQ_POST_SERIALIZE = 0x2 # cb expects Netbios frame 98 EV_REQ_PRE_SEND = 0x3 # cb expects a buffer to send 99 EV_REQ_POST_SEND = 0x4 # cb expects an integer of bytes sent 100 EV_RES_PRE_RECV = 0x5 # cb expects an integer of bytes to read 101 EV_RES_POST_RECV = 0x6 # cb expects a buffer that was read 102 EV_RES_PRE_DESERIALIZE = 0x7 # cb expects a complete netbios buffer 103 EV_RES_POST_DESERIALIZE = 0x8 # cb expects a Netbios frame

104 Events.import_items(globals())

105 106 -class Future(object):

107 """ 108 Result of an asynchronous operation. 109 110 Futures represent the result of an operation that has not yet completed. 111 In Pike, they are most commonly used to track SMB2 request/response pairs, 112 but they can be used for any asynchronous operation. 113 114 The result of a future can be waited for synchronously by simply calling 115 L{Future.result}, or a notification callback can be set with L{Future.then}. 116 117 Futures implement the context manager interface so that they can be used 118 as the context for a with block. If an exception is raised from the block, 119 it will automatically be set as the result of the future. 120 121 @ivar request: The request associated with the future, usually an SMB2 request frame. 122 @ivar response: The result of the future, usually an SMB2 response frame. 123 @ivar interim_response: The interim response, usually an SMB2 response frame. 124 @ivar traceback: The traceback of an exception result, if applicable. 125 """ 126

127 - def __init__(self, request=None):

128 """ 129 Initialize future. 130 131 @param request: The request associated with the response. 132 """ 133 self.request = request 134 self.interim_response = None 135 self.response = None 136 self.notify = None 137 self.traceback = None

138

139 - def complete(self, response, traceback=None):

140 """ 141 Completes the future with the given result. 142 143 Calling a future as a function is equivalent to calling this method. 144 145 @param response: The result of the future. 146 @param traceback: If response is an exception, an optional traceback 147 """ 148 self.response = response 149 self.traceback = traceback 150 if self.notify is not None: 151 self.notify(self)

152

153 - def interim(self, response):

154 """ 155 Set interim response. 156 157 @param response: The interim response. 158 """ 159 self.interim_response = response

160

161 - def wait(self, timeout=default_timeout):

162 """ 163 Wait for future result to become available. 164 165 @param timeout: The time in seconds before giving up and raising TimeoutError 166 """ 167 deadline = time.time() + timeout 168 while self.response is None: 169 now = time.time() 170 if now > deadline: 171 raise TimeoutError('Timed out after %s seconds' % timeout) 172 loop(timeout=deadline-now, count=1) 173 174 return self

175

176 - def wait_interim(self, timeout=default_timeout):

177 """ 178 Wait for interim response or actual result to become available. 179 180 @param timeout: The time in seconds before giving up and raising TimeoutError 181 """ 182 deadline = time.time() + timeout 183 while self.response is None and self.interim_response is None: 184 now = time.time() 185 if now > deadline: 186 raise TimeoutError('Timed out after %s seconds' % timeout) 187 loop(timeout=deadline-now, count=1) 188 189 return self

190

191 - def result(self, timeout=default_timeout):

192 """ 193 Return result of future. 194 195 If the result is not yet available, this function will wait for it. 196 If the result is an exception, this function will raise it instead of 197 returning it. 198 199 @param timeout: The time in seconds before giving up and raising TimeoutError 200 """ 201 self.wait(timeout) 202 203 if isinstance(self.response, BaseException): 204 traceback = self.traceback 205 self.traceback = None 206 raise self.response, None, traceback 207 else: 208 return self.response

209

210 - def then(self, notify):

211 """ 212 Set notification function. 213 214 @param notify: A function which will be invoked with this future as a parameter 215 when its result becomes available. If it is already available, 216 it will be called immediately. 217 """ 218 if self.response is not None: 219 notify(self) 220 else: 221 self.notify = notify

222

223 - def __enter__(self):

224 pass

225

226 - def __exit__(self, exc_type, exc_value, traceback):

227 if exc_type is not None: 228 self.complete(exc_value, traceback) 229 return True

230

231 - def __call__(self, *params, **kwparams):

232 self.complete(*params, **kwparams)

233

234 -class Client(object):

235 """ 236 Client 237 238 Maintains all state associated with an SMB2/3 client. 239 240 @type dialects: [number] 241 @ivar dialects: A list of supported dialects 242 @ivar capabilities: Capabilities flags 243 @ivar security_mode: Security mode flags 244 @ivar client_guid: Client GUID 245 @ivar channel_sequence: Current channel sequence number 246 """

247 - def __init__(self, 248 dialects=[smb2.DIALECT_SMB2_002, 249 smb2.DIALECT_SMB2_1, 250 smb2.DIALECT_SMB3_0, 251 smb2.DIALECT_SMB3_0_2, 252 smb2.DIALECT_SMB3_1_1], 253 capabilities=smb2.GlobalCaps(reduce(operator.or_, smb2.GlobalCaps.values())), 254 security_mode=smb2.SMB2_NEGOTIATE_SIGNING_ENABLED, 255 client_guid=None):

256 """ 257 Constructor. 258 259 @type dialects: [number] 260 @param dialects: A list of supported dialects. 261 @param capabilities: Client capabilities flags 262 @param security_mode: Client security mode flags 263 @param client_guid: Client GUID. If None, a new one will be generated at random. 264 """ 265 object.__init__(self) 266 267 if client_guid is None: 268 client_guid = array.array('B',map(random.randint, [0]*16, [255]*16)) 269 270 self.dialects = dialects 271 self.capabilities = capabilities 272 self.security_mode = security_mode 273 self.client_guid = client_guid 274 self.channel_sequence = 0 275 self.callbacks = {} 276 self._oplock_break_map = {} 277 self._lease_break_map = {} 278 self._oplock_break_queue = [] 279 self._lease_break_queue = [] 280 self._connections = [] 281 self._leases = {} 282 283 self.logger = logging.getLogger('pike')

284 285 @contextlib.contextmanager

286 - def callback(self, event, cb):

287 """ 288 Register a callback function for the context block, then unregister it 289 """ 290 self.register_callback(event, cb) 291 try: 292 yield 293 finally: 294 self.unregister_callback(event, cb)

295

296 - def register_callback(self, event, cb):

297 """ 298 Registers a callback function, cb for the given event. 299 When the event fires, cb will be called with the relevant top-level 300 Netbios frame as the single paramter. 301 """ 302 ev = Events(event) 303 if ev not in self.callbacks: 304 self.callbacks[ev] = [] 305 self.callbacks[ev].append(cb)

306

307 - def unregister_callback(self, event, cb):

308 """ 309 Unregisters a callback function, cb for the given event. 310 """ 311 ev = Events(event) 312 if ev not in self.callbacks: 313 return 314 if cb not in self.callbacks[ev]: 315 return 316 self.callbacks[ev].remove(cb)

317

318 - def connect(self, server, port=445):

319 """ 320 Create a connection. 321 322 Returns a new L{Connection} object connected to the given server and port. 323 324 @param server: The server to connect to. 325 @param port: The port to connect to. 326 """ 327 return self.connect_submit(server, port).result()

328

329 - def connect_submit(self, server, port=445):

330 """ 331 Create a connection. 332 333 Returns a new L{Future} object for the L{Connection} being established 334 asynchronously to the given server and port. 335 336 @param server: The server to connect to. 337 @param port: The port to connect to. 338 """ 339 return Connection(self, server, port).connection_future

340 341 # Do not use, may be removed. Use oplock_break_future.

342 - def next_oplock_break(self):

343 while len(self._oplock_break_queue) == 0: 344 loop(count=1) 345 return self._oplock_break_queue.pop()

346 347 # Do not use, may be removed. Use lease_break_future.

348 - def next_lease_break(self):

349 while len(self._lease_break_queue) == 0: 350 loop(count=1) 351 return self._lease_break_queue.pop()

352

353 - def oplock_break_future(self, file_id):

354 """ 355 Create future for oplock break. 356 357 Returns a L{Future} object which will be completed when 358 an oplock break occurs. The result will be the L{smb2.Smb2} frame 359 of the break notification packet. 360 361 @type file_id: (number, number) 362 @param file_id: The file ID of the oplocked file. 363 """ 364 365 future = Future(None) 366 367 for smb_res in self._oplock_break_queue[:]: 368 if smb_res[0].file_id == file_id: 369 future.complete(smb_res) 370 self._oplock_break_queue.remove(smb_res) 371 break 372 373 if future.response is None: 374 self._oplock_break_map[file_id] = future 375 376 return future

377

378 - def lease_break_future(self, lease_key):

379 """ 380 Create future for lease break. 381 382 Returns a L{Future} object which will be completed when 383 a lease break occurs. The result will be the L{smb2.Smb2} frame 384 of the break notification packet. 385 386 @param lease_key: The lease key for the lease. 387 """ 388 389 future = Future(None) 390 391 for smb_res in self._lease_break_queue[:]: 392 if smb_res[0].lease_key == lease_key: 393 future.complete(smb_res) 394 self._lease_break_queue.remove(smb_res) 395 break 396 397 if future.response is None: 398 self._lease_break_map[lease_key.tostring()] = future 399 400 return future

401

402 - def oplock_break(self, file_id):

403 """ 404 Wait for and return oplock break notification. 405 406 Equivalent to L{oplock_break_future}(file_id).result() 407 """ 408 409 return self.oplock_break_future(file_id).result()

410

411 - def lease_break(self, lease_key):

412 """ 413 Wait for and return lease break notification. 414 415 Equivalent to L{lease_break_future}(lease_key).result() 416 """ 417 418 return self.lease_break_future(lease_key).result()

419

420 - def lease(self, tree, lease_res):

421 """ 422 Create or look up lease object. 423 424 Returns a lease object based on a L{Tree} and a 425 L{smb2.LeaseResponse}. The lease object is created 426 if it does not already exist. 427 428 @param tree: The tree on which the lease request was issued. 429 @param lease_res: The lease create context response. 430 """ 431 432 lease_key = lease_res.lease_key.tostring() 433 if lease_key not in self._leases: 434 lease = Lease(tree) 435 self._leases[lease_key] = lease 436 else: 437 lease = self._leases[lease_key] 438 lease.ref() 439 440 lease.update(lease_res) 441 return lease

442 443 # Internal function to remove lease from table

444 - def dispose_lease(self, lease):

445 del self._leases[lease.lease_key.tostring()]

446

447 -class Connection(transport.Transport):

448 """ 449 Connection to server. 450 451 Represents a connection to a server and handles all socket operations 452 and request/response dispatch. 453 454 @type client: Client 455 @ivar client: The Client object associated with this connection. 456 @ivar server: The server name or address 457 @ivar port: The server port 458 """

459 - def __init__(self, client, server, port=445):

460 """ 461 Constructor. 462 463 This should generally not be used directly. Instead, 464 use L{Client.connect}(). 465 """ 466 super(Connection, self).__init__() 467 self._no_delay = True 468 self._in_buffer = array.array('B') 469 self._watermark = 4 470 self._out_buffer = None 471 self._next_mid = 0 472 self._mid_blacklist = set() 473 self._out_queue = [] 474 self._future_map = {} 475 self._sessions = {} 476 self._binding = None 477 self._binding_key = None 478 self._settings = {} 479 self._pre_auth_integrity_hash = array.array('B', "\0"*64) 480 self.callbacks = {} 481 self.connection_future = Future() 482 self.credits = 0 483 self.client = client 484 self.server = server 485 self.port = port 486 self.remote_addr = None 487 self.local_addr = None 488 self.verify_signature = True 489 490 self.error = None 491 self.traceback = None 492 493 for result in socket.getaddrinfo(server, port, 494 0, 495 socket.SOCK_STREAM, 496 socket.IPPROTO_TCP): 497 family, socktype, proto, canonname, sockaddr = result 498 break 499 self.create_socket(family, socktype) 500 self.connect(sockaddr)

501 502 @contextlib.contextmanager

503 - def callback(self, event, cb):

504 """ 505 Register a callback function for the context block, then unregister it 506 """ 507 self.register_callback(event, cb) 508 try: 509 yield 510 finally: 511 self.unregister_callback(event, cb)

512

513 - def register_callback(self, event, cb):

514 """ 515 Registers a callback function, cb for the given event. 516 When the event fires, cb will be called with the relevant top-level 517 Netbios frame as the single paramter. 518 """ 519 ev = Events(event) 520 if ev not in self.callbacks: 521 self.callbacks[ev] = [] 522 self.callbacks[ev].append(cb)

523

524 - def unregister_callback(self, event, cb):

525 """ 526 Unregisters a callback function, cb for the given event. 527 """ 528 ev = Events(event) 529 if ev not in self.callbacks: 530 return 531 if cb not in self.callbacks[ev]: 532 return 533 self.callbacks[ev].remove(cb)

534

535 - def process_callbacks(self, event, obj):

536 """ 537 Fire callbacks for the given event, passing obj as the parameter 538 539 Connection-specific callbacks will be fired first, followed by client 540 callbacks 541 """ 542 ev = Events(event) 543 all_callbacks = [self.callbacks] 544 if hasattr(self.client, "callbacks"): 545 all_callbacks.append(self.client.callbacks) 546 for callbacks in all_callbacks: 547 if ev not in callbacks: 548 continue 549 for cb in callbacks[ev]: 550 cb(obj)

551

552 - def smb3_pa_integrity(self, packet, data=None):

553 """ perform smb3 pre-auth integrity hash update if needed """ 554 if smb2.DIALECT_SMB3_1_1 not in self.client.dialects: 555 # hash only applies if client requests 3.1.1 556 return 557 neg_resp = getattr(self, "negotiate_response", None) 558 if (neg_resp is not None and 559 neg_resp.dialect_revision < smb2.DIALECT_SMB3_1_1): 560 # hash only applies if server negotiates 3.1.1 561 return 562 if packet[0].__class__ not in [smb2.NegotiateRequest, 563 smb2.NegotiateResponse, 564 smb2.SessionSetupRequest, 565 smb2.SessionSetupResponse]: 566 # hash only applies to pre-auth messages 567 return 568 if (packet[0].__class__ == smb2.SessionSetupResponse and 569 packet.status == ntstatus.STATUS_SUCCESS): 570 # last session setup doesn't count in hash 571 return 572 if data is None: 573 data = packet.serialize() 574 self._pre_auth_integrity_hash = digest.smb3_sha512( 575 self._pre_auth_integrity_hash + 576 data)

577

578 - def next_mid_range(self, length):

579 """ 580 multicredit requests must reserve 1 message id per credit charged. 581 the message id of the request should be the first id of the range. 582 """ 583 if length < 1: 584 length = 1 585 start_range = self._next_mid 586 while True: 587 r = set(range(start_range, start_range+length)) 588 if not r.intersection(self._mid_blacklist): 589 break 590 start_range += 1 591 self._next_mid = sorted(list(r))[-1] + 1 592 return start_range

593

594 - def next_mid(self):

595 return self.next_range(1)

596

597 - def reserve_mid(mid):

598 self._mid_blacklist.add(mid)

599

600 - def handle_connect(self):

601 self.client._connections.append(self) 602 with self.connection_future: 603 self.local_addr = self.socket.getsockname() 604 self.remote_addr = self.socket.getpeername() 605 606 self.client.logger.debug('connect: %s/%s -> %s/%s', 607 self.local_addr[0], self.local_addr[1], 608 self.remote_addr[0], self.remote_addr[1]) 609 self.connection_future(self)

610

611 - def handle_read(self):

612 # Try to read the next netbios frame 613 remaining = self._watermark - len(self._in_buffer) 614 self.process_callbacks(EV_RES_PRE_RECV, remaining) 615 data = array.array('B', self.recv(remaining)) 616 self.process_callbacks(EV_RES_POST_RECV, data) 617 self._in_buffer.extend(data) 618 avail = len(self._in_buffer) 619 if avail >= 4: 620 self._watermark = 4 + struct.unpack('>L', self._in_buffer[0:4])[0] 621 if avail == self._watermark: 622 nb = self.frame() 623 self.process_callbacks(EV_RES_PRE_DESERIALIZE, self._in_buffer) 624 nb.parse(self._in_buffer) 625 self._in_buffer = array.array('B') 626 self._watermark = 4 627 self._dispatch_incoming(nb)

628

629 - def handle_write(self):

630 # Try to write out more data 631 while self._out_buffer is None and len(self._out_queue): 632 self._out_buffer = self._prepare_outgoing() 633 while self._out_buffer is not None: 634 self.process_callbacks(EV_REQ_PRE_SEND, self._out_buffer) 635 sent = self.send(self._out_buffer) 636 del self._out_buffer[:sent] 637 if len(self._out_buffer) == 0: 638 self._out_buffer = None 639 self.process_callbacks(EV_REQ_POST_SEND, sent)

640

641 - def handle_close(self):

642 self.close()

643

644 - def handle_error(self):

645 (_,self.error,self.traceback) = sys.exc_info() 646 self.close()

647

648 - def close(self):

649 """ 650 Close connection. 651 652 This unceremoniously terminates the connection and fails all 653 outstanding requests with EOFError. 654 """ 655 # If there is no error, propagate EOFError 656 if self.error is None: 657 self.error = EOFError("close") 658 659 # if the connection hasn't been established, raise the error 660 if self.connection_future.response is None: 661 self.connection_future(self.error) 662 663 # otherwise, ignore this connection since it's not associated with its client 664 if self not in self.client._connections: 665 return 666 667 super(Connection, self).close() 668 669 if self.remote_addr is not None: 670 self.client.logger.debug("disconnect (%s/%s -> %s/%s): %s", 671 self.local_addr[0], self.local_addr[1], 672 self.remote_addr[0], self.remote_addr[1], 673 self.error) 674 675 self.client._connections.remove(self) 676 677 for future in self._out_queue: 678 future.complete(self.error, self.traceback) 679 del self._out_queue[:] 680 681 for future in self._future_map.itervalues(): 682 future.complete(self.error, self.traceback) 683 self._future_map.clear() 684 685 for session in self._sessions.values(): 686 session.delchannel(self) 687 688 self.traceback = None

689

690 - def _prepare_outgoing(self):

691 # Try to prepare an outgoing packet 692 693 # Grab an outgoing smb2 request 694 future = self._out_queue[0] 695 696 result = None 697 with future: 698 req = future.request 699 self.process_callbacks(EV_REQ_PRE_SERIALIZE, req.parent) 700 701 if req.credit_charge is None: 702 req.credit_charge = 0 703 for cmd in req: 704 if isinstance(cmd, smb2.ReadRequest) and cmd.length > 0: 705 # special handling, 1 credit per 64k 706 req.credit_charge, remainder = divmod(cmd.length, 2**16) 707 elif isinstance(cmd, smb2.WriteRequest) and cmd.buffer is not None: 708 # special handling, 1 credit per 64k 709 if cmd.length is None: 710 cmd.length = len(cmd.buffer) 711 req.credit_charge, remainder = divmod(cmd.length, 2**16) 712 else: 713 remainder = 1 # assume 1 credit per command 714 if remainder > 0: 715 req.credit_charge += 1 716 # do credit accounting based on our calculations (MS-SMB2 3.2.5.1) 717 self.credits -= req.credit_charge 718 719 if req.credit_request is None: 720 req.credit_request = default_credit_request 721 if req.credit_charge > req.credit_request: 722 req.credit_request = req.credit_charge # try not to fall behind 723 724 del self._out_queue[0] 725 726 # Assign message id 727 if req.message_id is None: 728 req.message_id = self.next_mid_range(req.credit_charge) 729 730 if req.is_last_child(): 731 # Last command in chain, ready to send packet 732 # TODO: move smb pa integrity to callback 733 self.smb3_pa_integrity(req) 734 result = req.parent.serialize() 735 self.process_callbacks(EV_REQ_POST_SERIALIZE, req.parent) 736 if trace: 737 self.client.logger.debug('send (%s/%s -> %s/%s): %s', 738 self.local_addr[0], self.local_addr[1], 739 self.remote_addr[0], self.remote_addr[1], 740 req.parent) 741 else: 742 self.client.logger.debug('send (%s/%s -> %s/%s): %s', 743 self.local_addr[0], self.local_addr[1], 744 self.remote_addr[0], self.remote_addr[1], 745 ', '.join(f[0].__class__.__name__ for f in req.parent)) 746 else: 747 # Not ready to send chain 748 result = None 749 750 # Move it to map for response waiters (but not cancel) 751 if not isinstance(req[0], smb2.Cancel): 752 self._future_map[req.message_id] = future 753 754 return result

755

756 - def _find_oplock_future(self, file_id):

757 if file_id in self.client._oplock_break_map: 758 return self.client._oplock_break_map.pop(file_id) 759 return None

760

761 - def _find_lease_future(self, lease_key):

762 lease_key = lease_key.tostring() 763 if lease_key in self.client._lease_break_map: 764 return self.client._lease_break_map.pop(lease_key) 765 return None

766

767 - def _dispatch_incoming(self, res):

768 if trace: 769 self.client.logger.debug('recv (%s/%s -> %s/%s): %s', 770 self.remote_addr[0], self.remote_addr[1], 771 self.local_addr[0], self.local_addr[1], 772 res) 773 else: 774 self.client.logger.debug('recv (%s/%s -> %s/%s): %s', 775 self.remote_addr[0], self.remote_addr[1], 776 self.local_addr[0], self.local_addr[1], 777 ', '.join(f[0].__class__.__name__ for f in res)) 778 self.process_callbacks(EV_RES_POST_DESERIALIZE, res) 779 for smb_res in res: 780 # TODO: move smb pa integrity and credit tracking to callbacks 781 self.smb3_pa_integrity(smb_res, smb_res.parent.buf[4:]) 782 self.credits += smb_res.credit_response 783 784 # Verify non-session-setup-response signatures 785 # session setup responses are verified in SessionSetupContext 786 if not isinstance(smb_res[0], smb2.SessionSetupResponse): 787 key = self.signing_key(smb_res.session_id) 788 if key and self.verify_signature: 789 smb_res.verify(self.signing_digest(), key) 790 791 if smb_res.message_id == smb2.UNSOLICITED_MESSAGE_ID: 792 if isinstance(smb_res[0], smb2.OplockBreakNotification): 793 future = self._find_oplock_future(smb_res[0].file_id) 794 if future: 795 future.complete(smb_res) 796 else: 797 self.client._oplock_break_queue.append(smb_res) 798 elif isinstance(smb_res[0], smb2.LeaseBreakNotification): 799 future = self._find_lease_future(smb_res[0].lease_key) 800 if future: 801 future.complete(smb_res) 802 else: 803 self.client._lease_break_queue.append(smb_res) 804 else: 805 raise core.BadPacket() 806 elif smb_res.message_id in self._future_map: 807 future = self._future_map[smb_res.message_id] 808 if smb_res.status == ntstatus.STATUS_PENDING: 809 future.interim(smb_res) 810 elif isinstance(smb_res[0], smb2.ErrorResponse) or \ 811 smb_res.status not in smb_res[0].allowed_status: 812 future.complete(ResponseError(smb_res)) 813 del self._future_map[smb_res.message_id] 814 else: 815 future.complete(smb_res) 816 del self._future_map[smb_res.message_id]

817

818 - def submit(self, req):

819 """ 820 Submit request. 821 822 Submits a L{netbios.Netbios} frame for sending. Returns 823 a list of L{Future} objects, one for each corresponding 824 L{smb2.Smb2} frame in the request. 825 """ 826 if self.error is not None: 827 raise self.error, None, self.traceback 828 futures = [] 829 for smb_req in req: 830 if isinstance(smb_req[0], smb2.Cancel): 831 # Find original future being canceled to return 832 if smb_req.async_id is not None: 833 # Cancel by async ID 834 future = filter(lambda f: f.interim_response.async_id == smb_req.async_id, self._future_map.itervalues())[0] 835 elif smb_req.message_id in self._future_map: 836 # Cancel by message id, already in future map 837 future = self._future_map[smb_req.message_id] 838 else: 839 # Cancel by message id, still in send queue 840 future = filter(lambda f: f.request.message_id == smb_req.message_id, self._out_queue)[0] 841 # Add fake future for cancel since cancel has no response 842 self._out_queue.append(Future(smb_req)) 843 futures.append(future) 844 else: 845 future = Future(smb_req) 846 self._out_queue.append(future) 847 futures.append(future) 848 849 # don't wait for the callback, send the data now 850 if self._no_delay: 851 self.handle_write() 852 return futures

853

854 - def transceive(self, req):

855 """ 856 Submit request and wait for responses. 857 858 Submits a L{netbios.Netbios} frame for sending. Waits for 859 and returns a list of L{smb2.Smb2} response objects, one for each 860 corresponding L{smb2.Smb2} frame in the request. 861 """ 862 return map(Future.result, self.submit(req))

863

864 - def negotiate_request(self, hash_algorithms=None, salt=None, ciphers=None):

865 smb_req = self.request() 866 smb_req.credit_charge = 0 # negotiate requests are free 867 neg_req = smb2.NegotiateRequest(smb_req) 868 869 neg_req.dialects = self.client.dialects 870 neg_req.security_mode = self.client.security_mode 871 neg_req.capabilities = self.client.capabilities 872 neg_req.client_guid = self.client.client_guid 873 874 if smb2.DIALECT_SMB3_1_1 in neg_req.dialects: 875 if ciphers is None: 876 ciphers = [crypto.SMB2_AES_128_GCM, 877 crypto.SMB2_AES_128_CCM] 878 if ciphers: 879 encryption_req = crypto.EncryptionCapabilitiesRequest(neg_req) 880 encryption_req.ciphers = ciphers 881 882 preauth_integrity_req = smb2.PreauthIntegrityCapabilitiesRequest(neg_req) 883 if hash_algorithms is None: 884 hash_algorithms = [smb2.SMB2_SHA_512] 885 preauth_integrity_req.hash_algorithms = hash_algorithms 886 if salt is not None: 887 preauth_integrity_req.salt = salt 888 else: 889 preauth_integrity_req.salt = array.array('B', 890 map(random.randint, [0]*32, [255]*32)) 891 return neg_req

892

893 - def negotiate_submit(self, negotiate_request):

894 negotiate_future = self.submit(negotiate_request.parent.parent)[0] 895 def assign_response(f): 896 self.negotiate_response = f.result()[0]

897 negotiate_future.then(assign_response) 898 return negotiate_future

899

900 - def negotiate(self, hash_algorithms=None, salt=None, ciphers=None):

901 """ 902 Perform dialect negotiation. 903 904 This must be performed before setting up a session with 905 L{Connection.session_setup}(). 906 """ 907 self.negotiate_submit( 908 self.negotiate_request( 909 hash_algorithms, 910 salt, 911 ciphers 912 )).result() 913 return self

914

915 - class SessionSetupContext(object):

916 - def __init__(self, conn, creds=None, bind=None, resume=None, 917 ntlm_version=None):

918 assert conn.negotiate_response is not None 919 920 self.conn = conn 921 self.dialect_revision = conn.negotiate_response.dialect_revision 922 self.bind = bind 923 self.resume = resume 924 925 if creds and auth.ntlm is not None: 926 self.auth = auth.NtlmProvider(conn, creds) 927 if ntlm_version is not None: 928 self.auth.authenticator.ntlm_version = ntlm_version 929 elif auth.kerberos is not None: 930 self.auth = auth.KerberosProvider(conn, creds) 931 else: 932 raise ImportError("Neither ntlm nor kerberos authentication " 933 "methods are available") 934 935 self._settings = {} 936 self.prev_session_id = 0 937 self.session_id = 0 938 self.requests = [] 939 self.responses = [] 940 self.session_future = Future() 941 self.interim_future = None 942 943 if bind: 944 assert conn.negotiate_response.dialect_revision >= 0x300 945 self.session_id = bind.session_id 946 conn._binding = bind 947 # assume the signing key from the previous session 948 conn._binding_key = bind.first_channel().signing_key 949 elif resume: 950 assert conn.negotiate_response.dialect_revision >= 0x300 951 self.prev_session_id = resume.session_id

952

953 - def let(self, **kwargs):

954 return core.Let(self, kwargs)

955

956 - def derive_signing_key(self, session_key=None, context=None):

957 if session_key is None: 958 session_key = self.session_key 959 if self.dialect_revision >= smb2.DIALECT_SMB3_1_1: 960 if context is None: 961 context = self.conn._pre_auth_integrity_hash 962 return digest.derive_key( 963 session_key, 964 'SMBSigningKey', 965 context)[:16] 966 elif self.dialect_revision >= smb2.DIALECT_SMB3_0: 967 if context is None: 968 context = 'SmbSign\0' 969 return digest.derive_key(session_key, 'SMB2AESCMAC', context)[:16] 970 else: 971 return session_key

972

973 - def derive_encryption_keys(self, session_key=None, context=None):

974 if self.dialect_revision >= smb2.DIALECT_SMB3_1_1: 975 if context is None: 976 context = self.conn._pre_auth_integrity_hash 977 for nctx in self.conn.negotiate_response: 978 if isinstance(nctx, crypto.EncryptionCapabilitiesResponse): 979 try: 980 return crypto.EncryptionContext( 981 crypto.CryptoKeys311( 982 self.session_key, 983 context), 984 nctx.ciphers) 985 except crypto.CipherMismatch: 986 pass 987 elif self.dialect_revision >= smb2.DIALECT_SMB3_0: 988 if self.conn.negotiate_response.capabilities & smb2.SMB2_GLOBAL_CAP_ENCRYPTION: 989 return crypto.EncryptionContext( 990 crypto.CryptoKeys300(self.session_key), 991 [crypto.SMB2_AES_128_CCM])

992

993 - def _send_session_setup(self, sec_buf):

994 smb_req = self.conn.request() 995 session_req = smb2.SessionSetupRequest(smb_req) 996 997 smb_req.session_id = self.session_id 998 session_req.previous_session_id = self.prev_session_id 999 session_req.security_mode = smb2.SMB2_NEGOTIATE_SIGNING_ENABLED 1000 session_req.security_buffer = sec_buf 1001 if self.bind: 1002 smb_req.flags = smb2.SMB2_FLAGS_SIGNED 1003 session_req.flags = smb2.SMB2_SESSION_FLAG_BINDING 1004 1005 for (attr,value) in self._settings.iteritems(): 1006 setattr(session_req, attr, value) 1007 1008 self.requests.append(smb_req) 1009 return self.conn.submit(smb_req.parent)[0]

1010

1011 - def _finish(self, smb_res):

1012 sec_buf = smb_res[0].security_buffer 1013 out_buf, self.session_key = self.auth.step(sec_buf) 1014 signing_key = self.derive_signing_key() 1015 encryption_context = self.derive_encryption_keys() 1016 1017 # Verify final signature 1018 smb_res.verify(self.conn.signing_digest(), signing_key) 1019 1020 if self.bind: 1021 self.conn._binding = None 1022 self.conn._binding_key = None 1023 session = self.bind 1024 else: 1025 session = Session(self.conn.client, 1026 self.session_id, 1027 self.session_key, 1028 encryption_context, 1029 smb_res) 1030 session.user = self.auth.username() 1031 1032 return session.addchannel(self.conn, signing_key)

1033

1034 - def __iter__(self):

1035 return self

1036

1037 - def submit(self, f=None):

1038 """ 1039 Submit rounds of SessionSetupRequests 1040 1041 Returns a L{Future} object, for the L{Channel} object 1042 """ 1043 try: 1044 res = self.next() 1045 res.then(self.submit) 1046 except StopIteration: 1047 pass 1048 return self.session_future

1049

1050 - def next(self):

1051 with self.session_future: 1052 res = self._process() 1053 if res is not None: 1054 return res 1055 raise StopIteration()

1056

1057 - def _process(self):

1058 out_buf = None 1059 if not self.interim_future and not self.responses: 1060 # send the initial request 1061 out_buf, self.session_key = self.auth.step( 1062 self.conn.negotiate_response.security_buffer) 1063 1064 elif self.interim_future: 1065 smb_res = self.interim_future.result() 1066 self.interim_future = None 1067 self.responses.append(smb_res) 1068 self.session_id = smb_res.session_id 1069 1070 if smb_res.status == ntstatus.STATUS_SUCCESS: 1071 # session is established 1072 with self.session_future: 1073 self.session_future(self._finish(smb_res)) 1074 return self.session_future 1075 else: 1076 # process interim request 1077 session_res = smb_res[0] 1078 if self.bind: 1079 # Need to verify intermediate signatures 1080 smb_res.verify(self.conn.signing_digest(), 1081 self.conn._binding_key) 1082 out_buf, self.session_key = self.auth.step( 1083 session_res.security_buffer) 1084 if out_buf: 1085 # submit additional requests if necessary 1086 self.interim_future = self._send_session_setup(out_buf) 1087 return self.interim_future

1088

1089 - def session_setup(self, creds=None, bind=None, resume=None):

1090 """ 1091 Establish a session. 1092 1093 Establishes a session, performing GSS rounds as necessary. Returns 1094 a L{Channel} object which can be used for further requests on the given 1095 connection and session. 1096 1097 @type creds: str 1098 @param creds: A set of credentials of the form '<domain>\<user>%<password>'. 1099 If specified, NTLM authentication will be used. If None, 1100 Kerberos authentication will be attempted. 1101 @type bind: L{Session} 1102 @param bind: An existing session to bind. 1103 @type resume: L{Session} 1104 @param resume: An previous session to resume. 1105 """ 1106 session_context = self.SessionSetupContext(self, creds, bind, resume) 1107 return session_context.submit().result()

1108 1109 # Return a fresh netbios frame with connection as context

1110 - def frame(self):

1111 return netbios.Netbios(context=self)

1112 1113 # Return a fresh smb2 frame with connection as context 1114 # Put it in a netbios frame automatically if none given

1115 - def request(self, parent=None):

1116 if parent is None: 1117 parent = self.frame() 1118 req = smb2.Smb2(parent, context=self) 1119 req.channel_sequence = self.client.channel_sequence 1120 1121 for (attr,value) in self._settings.iteritems(): 1122 setattr(req, attr, value) 1123 1124 return req

1125

1126 - def let(self, **kwargs):

1127 return core.Let(self, kwargs)

1128 1129 # 1130 # SMB2 context upcalls 1131 #

1132 - def session(self, session_id):

1133 return self._sessions.get(session_id, None)

1134

1135 - def signing_key(self, session_id):

1136 if session_id in self._sessions: 1137 session = self._sessions[session_id] 1138 channel = session._channels[id(self)] 1139 return channel.signing_key 1140 elif self._binding and self._binding.session_id == session_id: 1141 return self._binding_key

1142

1143 - def encryption_context(self, session_id):

1144 if session_id in self._sessions: 1145 session = self._sessions[session_id] 1146 return session.encryption_context

1147

1148 - def signing_digest(self):

1149 assert self.negotiate_response is not None 1150 if self.negotiate_response.dialect_revision >= smb2.DIALECT_SMB3_0: 1151 return digest.aes128_cmac 1152 else: 1153 return digest.sha256_hmac

1154

1155 - def get_request(self, message_id):

1156 if message_id in self._future_map: 1157 return self._future_map[message_id].request 1158 else: 1159 return None

1160

1161 -class Session(object):

1162 - def __init__(self, client, session_id, session_key, 1163 encryption_context, smb_res):

1164 object.__init__(self) 1165 self.client = client 1166 self.session_id = session_id 1167 self.session_key = session_key 1168 self.encryption_context = encryption_context 1169 self.encrypt_data = False 1170 if smb_res[0].session_flags & smb2.SMB2_SESSION_FLAG_ENCRYPT_DATA and \ 1171 self.encryption_context is not None: 1172 self.encrypt_data = True 1173 self._channels = {} 1174 self._trees = {} 1175 self.user = None

1176

1177 - def addchannel(self, conn, signing_key):

1178 channel = Channel(conn, self, signing_key) 1179 self._channels[id(conn)] = channel 1180 conn._sessions[self.session_id] = self 1181 return channel

1182

1183 - def delchannel(self, conn):

1184 del conn._sessions[self.session_id] 1185 del self._channels[id(conn)]

1186

1187 - def first_channel(self):

1188 return self._channels.itervalues().next()

1189

1190 - def tree(self, tree_id):

1191 return self._trees.get(tree_id, None)

1192

1193 -class Channel(object):

1194 - def __init__(self, connection, session, signing_key):

1195 object.__init__(self) 1196 self.connection = connection 1197 self.session = session 1198 self.signing_key = signing_key

1199

1200 - def cancel_request(self, future):

1201 if (future.response is not None): 1202 raise StateError("Cannot cancel completed request") 1203 1204 smb_req = self.request() 1205 cancel_req = smb2.Cancel(smb_req) 1206 1207 # Don't bother trying to sign cancel 1208 smb_req.flags &= ~smb2.SMB2_FLAGS_SIGNED 1209 1210 # Use async id to cancel if applicable: 1211 if future.interim_response is not None: 1212 smb_req.async_id = future.interim_response.async_id 1213 smb_req.tree_id = None 1214 smb_req.flags |= smb2.SMB2_FLAGS_ASYNC_COMMAND 1215 smb_req.message_id = 0 1216 else: 1217 smb_req.message_id = future.request.message_id 1218 1219 return cancel_req

1220

1221 - def cancel(self, future):

1222 cancel_req = self.cancel_request(future) 1223 return self.connection.submit(cancel_req.parent.parent)[0]

1224

1225 - def tree_connect_request(self, path):

1226 smb_req = self.request() 1227 if self.connection.negotiate_response.dialect_revision >= smb2.DIALECT_SMB3_1_1: 1228 smb_req.flags |= smb2.SMB2_FLAGS_SIGNED 1229 tree_req = smb2.TreeConnectRequest(smb_req) 1230 tree_req.path = "\\\\" + self.connection.server + "\\" + path 1231 return tree_req

1232

1233 - def tree_connect_submit(self, tree_req):

1234 tree_future = Future() 1235 resp_future = self.connection.submit(tree_req.parent.parent)[0] 1236 resp_future.then(lambda f: tree_future.complete(Tree(self.session, 1237 tree_req.path, 1238 f.result()))) 1239 return tree_future

1240

1241 - def tree_connect(self, path):

1242 return self.tree_connect_submit( 1243 self.tree_connect_request( 1244 path)).result()

1245

1246 - def tree_disconnect_request(self, tree):

1247 smb_req = self.request(obj=tree) 1248 tree_req = smb2.TreeDisconnectRequest(smb_req) 1249 return tree_req

1250

1251 - def tree_disconnect(self, tree):

1252 return self.connection.transceive( 1253 self.tree_disconnect_request(tree).parent.parent)[0]

1254

1255 - def logoff_request(self):

1256 smb_req = self.request() 1257 logoff_req = smb2.LogoffRequest(smb_req) 1258 return logoff_req

1259

1260 - def logoff_submit(self, logoff_req):

1261 def logoff_finish(f): 1262 for channel in self.session._channels.itervalues(): 1263 del channel.connection._sessions[self.session.session_id]

1264 logoff_future = self.connection.submit(logoff_req.parent.parent)[0] 1265 logoff_future.then(logoff_finish) 1266 return logoff_future

1267

1268 - def logoff(self):

1269 return self.logoff_submit( 1270 self.logoff_request()).result()

1271

1272 - def create_request( 1273 self, 1274 tree, 1275 path, 1276 access=smb2.GENERIC_READ | smb2.GENERIC_WRITE, 1277 attributes=smb2.FILE_ATTRIBUTE_NORMAL, 1278 share=0, 1279 disposition=smb2.FILE_OPEN_IF, 1280 options=0, 1281 maximal_access=None, 1282 oplock_level=smb2.SMB2_OPLOCK_LEVEL_NONE, 1283 lease_key=None, 1284 lease_state=None, 1285 durable=False, 1286 persistent=False, 1287 create_guid=None, 1288 app_instance_id=None, 1289 query_on_disk_id=False, 1290 extended_attributes=None, 1291 timewarp=None):

1292 1293 prev_open = None 1294 1295 smb_req = self.request(obj=tree) 1296 create_req = smb2.CreateRequest(smb_req) 1297 1298 create_req.name = path 1299 create_req.desired_access = access 1300 create_req.file_attributes = attributes 1301 create_req.share_access = share 1302 create_req.create_disposition = disposition 1303 create_req.create_options = options 1304 create_req.requested_oplock_level = oplock_level 1305 1306 if maximal_access: 1307 max_req = smb2.MaximalAccessRequest(create_req) 1308 if maximal_access is not True: 1309 max_req.timestamp = maximal_access 1310 1311 if oplock_level == smb2.SMB2_OPLOCK_LEVEL_LEASE: 1312 lease_req = smb2.LeaseRequest(create_req) 1313 lease_req.lease_key = lease_key 1314 lease_req.lease_state = lease_state 1315 1316 if isinstance(durable, Open): 1317 prev_open = durable 1318 if durable.durable_timeout is None: 1319 durable_req = smb2.DurableHandleReconnectRequest(create_req) 1320 durable_req.file_id = durable.file_id 1321 else: 1322 durable_req = smb2.DurableHandleReconnectV2Request(create_req) 1323 durable_req.file_id = durable.file_id 1324 durable_req.create_guid = durable.create_guid 1325 durable_req.flags = durable.durable_flags 1326 elif durable is True: 1327 durable_req = smb2.DurableHandleRequest(create_req) 1328 elif durable is not False: 1329 durable_req = smb2.DurableHandleV2Request(create_req) 1330 durable_req.timeout = durable 1331 if persistent: 1332 durable_req.flags = smb2.SMB2_DHANDLE_FLAG_PERSISTENT 1333 if create_guid is None: 1334 create_guid = array.array('B',map(random.randint, [0]*16, [255]*16)) 1335 durable_req.create_guid = create_guid 1336 1337 if app_instance_id: 1338 app_instance_id_req = smb2.AppInstanceIdRequest(create_req) 1339 app_instance_id_req.app_instance_id = app_instance_id 1340 1341 if query_on_disk_id: 1342 query_on_disk_id_req = smb2.QueryOnDiskIDRequest(create_req) 1343 1344 if extended_attributes: 1345 ext_attr_len = len(extended_attributes.keys()) 1346 for name, value in extended_attributes.iteritems(): 1347 ext_attr = smb2.ExtendedAttributeRequest(create_req) 1348 if ext_attr_len == 1: 1349 next_entry_offset = 0 1350 else: 1351 next_entry_offset = 10 + len(name) + len(value) 1352 ext_attr.next_entry_offset = next_entry_offset 1353 ext_attr.ea_name = name 1354 ext_attr.ea_name_length = len(name) 1355 ext_attr.ea_value = value 1356 ext_attr.ea_value_length = len(value) 1357 ext_attr_len = ext_attr_len - 1 1358 1359 if timewarp: 1360 timewarp_req = smb2.TimewarpTokenRequest(create_req) 1361 timewarp_req.timestamp = nttime.NtTime(timewarp) 1362 1363 open_future = Future(None) 1364 def finish(f): 1365 with open_future: open_future( 1366 Open( 1367 tree, 1368 f.result(), 1369 create_guid=create_guid, 1370 prev=prev_open))

1371 create_req.open_future = open_future 1372 create_req.finish = finish 1373 1374 return create_req 1375

1376 - def create_submit(self, create_req):

1377 open_future = create_req.open_future 1378 open_future.request_future = self.connection.submit( 1379 create_req.parent.parent)[0] 1380 open_future.request_future.then(create_req.finish) 1381 1382 return open_future

1383

1384 - def create( 1385 self, 1386 tree, 1387 path, 1388 access=smb2.GENERIC_READ | smb2.GENERIC_WRITE, 1389 attributes=smb2.FILE_ATTRIBUTE_NORMAL, 1390 share=0, 1391 disposition=smb2.FILE_OPEN_IF, 1392 options=0, 1393 maximal_access=None, 1394 oplock_level=smb2.SMB2_OPLOCK_LEVEL_NONE, 1395 lease_key=None, 1396 lease_state=None, 1397 durable=False, 1398 persistent=False, 1399 create_guid=None, 1400 app_instance_id=None, 1401 query_on_disk_id=False, 1402 extended_attributes=None, 1403 timewarp=None):

1404 return self.create_submit(self.create_request( 1405 tree, 1406 path, 1407 access, 1408 attributes, 1409 share, 1410 disposition, 1411 options, 1412 maximal_access, 1413 oplock_level, 1414 lease_key, 1415 lease_state, 1416 durable, 1417 persistent, 1418 create_guid, 1419 app_instance_id, 1420 query_on_disk_id, 1421 extended_attributes, 1422 timewarp))

1423

1424 - def close_request(self, handle):

1425 smb_req = self.request(obj=handle) 1426 close_req = smb2.CloseRequest(smb_req) 1427 1428 close_req.file_id = handle.file_id 1429 close_req.handle = handle 1430 return close_req

1431

1432 - def close_submit(self, close_req):

1433 resp_future = self.connection.submit(close_req.parent.parent)[0] 1434 resp_future.then(close_req.handle.dispose()) 1435 return resp_future

1436

1437 - def close(self, handle):

1438 return self.close_submit( 1439 self.close_request(handle)).result()

1440

1441 - def query_directory_request( 1442 self, 1443 handle, 1444 file_information_class=smb2.FILE_DIRECTORY_INFORMATION, 1445 flags=0, 1446 file_index=0, 1447 file_name='*', 1448 output_buffer_length=8192):

1449 smb_req = self.request(obj=handle) 1450 enum_req = smb2.QueryDirectoryRequest(smb_req) 1451 enum_req.file_id = handle.file_id 1452 enum_req.file_name = file_name 1453 enum_req.output_buffer_length = output_buffer_length 1454 enum_req.file_information_class = file_information_class 1455 enum_req.flags = flags 1456 enum_req.file_index = file_index 1457 return enum_req

1458

1459 - def query_directory(self, 1460 handle, 1461 file_information_class=smb2.FILE_DIRECTORY_INFORMATION, 1462 flags=0, 1463 file_index=0, 1464 file_name='*', 1465 output_buffer_length=8192):

1466 return self.connection.transceive( 1467 self.query_directory_request( 1468 handle, 1469 file_information_class, 1470 flags, 1471 file_index, 1472 file_name, 1473 output_buffer_length).parent.parent)[0][0]

1474

1475 - def enum_directory(self, 1476 handle, 1477 file_information_class=smb2.FILE_DIRECTORY_INFORMATION, 1478 file_name = '*', 1479 output_buffer_length=8192):

1480 while True: 1481 try: 1482 for info in self.query_directory(handle, 1483 file_information_class=file_information_class, 1484 file_name=file_name, 1485 output_buffer_length=output_buffer_length): 1486 yield info 1487 except ResponseError as e: 1488 if e.response.status == ntstatus.STATUS_NO_MORE_FILES: 1489 return 1490 else: 1491 raise

1492

1493 - def query_file_info_request( 1494 self, 1495 create_res, 1496 file_information_class=smb2.FILE_BASIC_INFORMATION, 1497 info_type=smb2.SMB2_0_INFO_FILE, 1498 output_buffer_length=4096, 1499 additional_information=None):

1500 smb_req = self.request(obj=create_res) 1501 query_req = smb2.QueryInfoRequest(smb_req) 1502 1503 query_req.info_type = info_type 1504 query_req.file_information_class = file_information_class 1505 query_req.file_id = create_res.file_id 1506 query_req.output_buffer_length = output_buffer_length 1507 if additional_information: 1508 query_req.additional_information = additional_information 1509 return query_req

1510

1511 - def query_file_info(self, 1512 create_res, 1513 file_information_class=smb2.FILE_BASIC_INFORMATION, 1514 info_type=smb2.SMB2_0_INFO_FILE, 1515 output_buffer_length=4096, 1516 additional_information=None):

1517 return self.connection.transceive( 1518 self.query_file_info_request( 1519 create_res, 1520 file_information_class, 1521 info_type, 1522 output_buffer_length, 1523 additional_information).parent.parent)[0][0][0]

1524

1525 - def set_file_info_request( 1526 self, 1527 handle, 1528 file_information_class=smb2.FILE_BASIC_INFORMATION, 1529 info_type=smb2.SMB2_0_INFO_FILE, 1530 input_buffer_length=4096, 1531 additional_information=None):

1532 smb_req = self.request(obj=handle) 1533 set_req = smb2.SetInfoRequest(smb_req) 1534 set_req.file_id = handle.file_id 1535 set_req.file_information_class = file_information_class 1536 set_req.info_type = info_type 1537 set_req.input_buffer_length = input_buffer_length 1538 if additional_information: 1539 set_req.additional_information = additional_information 1540 return set_req

1541 1542 @contextlib.contextmanager

1543 - def set_file_info(self, handle, cls):

1544 info_type = file_information_class = None 1545 if hasattr(cls, "info_type"): 1546 info_type = cls.info_type 1547 if hasattr(cls, "file_information_class"): 1548 file_information_class = cls.file_information_class 1549 set_req = self.set_file_info_request( 1550 handle, 1551 file_information_class, 1552 info_type) 1553 yield cls(set_req) 1554 self.connection.transceive(set_req.parent.parent)[0]

1555

1556 - def change_notify_request( 1557 self, 1558 handle, 1559 completion_filter=smb2.SMB2_NOTIFY_CHANGE_CREATION, 1560 flags=0, 1561 buffer_length=4096):

1562 smb_req = self.request(obj=handle) 1563 cnotify_req = smb2.ChangeNotifyRequest(smb_req) 1564 cnotify_req.file_id = handle.file_id 1565 cnotify_req.buffer_length = buffer_length 1566 cnotify_req.flags = flags 1567 return cnotify_req

1568

1569 - def change_notify( 1570 self, 1571 handle, 1572 completion_filter=smb2.SMB2_NOTIFY_CHANGE_CREATION, 1573 flags=0, 1574 buffer_length=4096):

1575 return self.connection.submit( 1576 self.change_notify_request( 1577 handle, 1578 completion_filter, 1579 flags, 1580 buffer_length=4096).parent.parent)[0][0]

1581 1582 # Send an echo request and get a response

1583 - def echo(self):

1584 # Create request structure 1585 smb_req = self.request() 1586 # Make the request struct have an ECHO_REQUEST 1587 enum_req = smb2.EchoRequest(smb_req) 1588 # Get response first [0] = first response, 2nd [0] = echo response 1589 # frame 1590 self.connection.transceive(smb_req.parent)[0][0]

1591

1592 - def flush_request(self, file):

1593 smb_req = self.request(obj=file) 1594 flush_req = smb2.FlushRequest(smb_req) 1595 flush_req.file_id = file.file_id 1596 return flush_req

1597

1598 - def flush(self, file):

1599 self.connection.transceive(self.flush_request(file).parent.parent)

1600

1601 - def read_request( 1602 self, 1603 file, 1604 length, 1605 offset, 1606 minimum_count=0, 1607 remaining_bytes=0):

1608 smb_req = self.request(obj=file) 1609 read_req = smb2.ReadRequest(smb_req) 1610 1611 read_req.length = length 1612 read_req.offset = offset 1613 read_req.minimum_count = minimum_count 1614 read_req.remaining_bytes = remaining_bytes 1615 read_req.file_id = file.file_id 1616 return read_req

1617

1618 - def read( 1619 self, 1620 file, 1621 length, 1622 offset, 1623 minimum_count=0, 1624 remaining_bytes=0):

1625 return self.connection.transceive( 1626 self.read_request( 1627 file, 1628 length, 1629 offset, 1630 minimum_count, 1631 remaining_bytes).parent.parent)[0][0].data

1632

1633 - def write_request( 1634 self, 1635 file, 1636 offset, 1637 buffer=None, 1638 remaining_bytes=0, 1639 flags=0):

1640 smb_req = self.request(obj=file) 1641 write_req = smb2.WriteRequest(smb_req) 1642 1643 write_req.offset = offset 1644 write_req.file_id = file.file_id 1645 write_req.buffer = buffer 1646 write_req.remaining_bytes = remaining_bytes 1647 write_req.flags = flags 1648 return write_req

1649

1650 - def write(self, 1651 file, 1652 offset, 1653 buffer=None, 1654 remaining_bytes=0, 1655 flags=0):

1656 smb_res = self.connection.transceive( 1657 self.write_request( 1658 file, 1659 offset, 1660 buffer, 1661 remaining_bytes, 1662 flags).parent.parent) 1663 1664 return smb_res[0][0].count

1665

1666 - def lock_request(self, handle, locks, sequence=0):

1667 """ 1668 @param locks: A list of lock tuples, each of which consists of (offset, length, flags). 1669 """ 1670 smb_req = self.request(obj=handle) 1671 lock_req = smb2.LockRequest(smb_req) 1672 1673 lock_req.file_id = handle.file_id 1674 lock_req.locks = locks 1675 lock_req.lock_sequence = sequence 1676 return lock_req

1677

1678 - def lock(self, handle, locks, sequence=0):

1679 """ 1680 @param locks: A list of lock tuples, each of which consists of (offset, length, flags). 1681 """ 1682 return self.connection.submit( 1683 self.lock_request( 1684 handle, 1685 locks, 1686 sequence).parent.parent)[0]

1687

1688 - def validate_negotiate_info(self, tree):

1689 smb_req = self.request(obj=tree) 1690 ioctl_req = smb2.IoctlRequest(smb_req) 1691 vni_req = smb2.ValidateNegotiateInfoRequest(ioctl_req) 1692 client = self.session.client 1693 1694 # Validate negotiate must always be signed 1695 smb_req.flags |= smb2.SMB2_FLAGS_SIGNED 1696 ioctl_req.flags = smb2.SMB2_0_IOCTL_IS_FSCTL 1697 vni_req.capabilities = client.capabilities 1698 vni_req.client_guid = client.client_guid 1699 vni_req.security_mode = client.security_mode 1700 vni_req.dialects = client.dialects 1701 1702 res = self.connection.transceive(smb_req.parent)[0] 1703 1704 return res

1705

1706 - def resume_key(self, file):

1707 smb_req = self.request(obj=file.tree) 1708 ioctl_req = smb2.IoctlRequest(smb_req) 1709 resumekey_req = smb2.RequestResumeKeyRequest(ioctl_req) 1710 1711 ioctl_req.file_id = file.file_id 1712 ioctl_req.flags |= smb2.SMB2_0_IOCTL_IS_FSCTL 1713 1714 return self.connection.transceive(smb_req.parent)[0]

1715

1716 - def copychunk_request(self, source_file, target_file, chunks):

1717 """ 1718 @param source_file: L{Open} 1719 @param target_file: L{Open} 1720 @param chunks: sequence of tuples (source_offset, target_offset, length) 1721 """ 1722 resume_key = self.resume_key(source_file)[0][0].resume_key 1723 1724 smb_req = self.request(obj=target_file.tree) 1725 ioctl_req = smb2.IoctlRequest(smb_req) 1726 copychunk_req = smb2.CopyChunkCopyRequest(ioctl_req) 1727 1728 ioctl_req.max_output_response = 16384 1729 ioctl_req.file_id = target_file.file_id 1730 ioctl_req.flags |= smb2.SMB2_0_IOCTL_IS_FSCTL 1731 copychunk_req.source_key = resume_key 1732 copychunk_req.chunk_count = len(chunks) 1733 1734 for source_offset, target_offset, length in chunks: 1735 chunk = smb2.CopyChunk(copychunk_req) 1736 chunk.source_offset = source_offset 1737 chunk.target_offset = target_offset 1738 chunk.length = length 1739 return ioctl_req

1740

1741 - def copychunk(self, source_file, target_file, chunks):

1742 """ 1743 @param source_file: L{Open} 1744 @param target_file: L{Open} 1745 @param chunks: sequence of tuples (source_offset, target_offset, length) 1746 """ 1747 return self.connection.transceive( 1748 self.copychunk_request( 1749 source_file, 1750 target_file, 1751 chunks).parent.parent)[0]

1752

1753 - def set_symlink_request(self, file, target_name, flags):

1754 smb_req = self.request(obj=file.tree) 1755 ioctl_req = smb2.IoctlRequest(smb_req) 1756 set_reparse_req = smb2.SetReparsePointRequest(ioctl_req) 1757 symlink_buffer = smb2.SymbolicLinkReparseBuffer(set_reparse_req) 1758 1759 ioctl_req.max_output_response = 0 1760 ioctl_req.file_id = file.file_id 1761 ioctl_req.flags |= smb2.SMB2_0_IOCTL_IS_FSCTL 1762 symlink_buffer.substitute_name = target_name 1763 symlink_buffer.flags = flags 1764 return ioctl_req

1765

1766 - def set_symlink(self, file, target_name, flags):

1767 return self.connection.transceive( 1768 self.set_symlink_request( 1769 file, 1770 target_name, 1771 flags).parent.parent)[0]

1772

1773 - def get_symlink_request(self, file):

1774 smb_req = self.request(obj=file.tree) 1775 ioctl_req = smb2.IoctlRequest(smb_req) 1776 set_reparse_req = smb2.GetReparsePointRequest(ioctl_req) 1777 1778 ioctl_req.file_id = file.file_id 1779 ioctl_req.flags |= smb2.SMB2_0_IOCTL_IS_FSCTL 1780 return ioctl_req

1781

1782 - def get_symlink(self, file):

1783 return self.connection.transceive( 1784 self.get_symlink_request(file).parent.parent)[0]

1785

1786 - def enumerate_snapshots_request(self, fh, max_output_response=16384):

1787 smb_req = self.request(obj=fh.tree) 1788 ioctl_req = smb2.IoctlRequest(smb_req) 1789 ioctl_req.max_output_response = max_output_response 1790 ioctl_req.file_id = fh.file_id 1791 ioctl_req.flags |= smb2.SMB2_0_IOCTL_IS_FSCTL 1792 enum_req = smb2.EnumerateSnapshotsRequest(ioctl_req) 1793 return enum_req

1794

1795 - def enumerate_snapshots(self, fh):

1796 return self.connection.transceive( 1797 self.enumerate_snapshots_request(fh).parent.parent.parent)[0]

1798

1799 - def enumerate_snapshots_list(self, fh):

1800 return self.enumerate_snapshots(fh)[0][0].snapshots

1801

1802 - def lease_break_acknowledgement(self, tree, notify):

1803 """ 1804 @param tree: L{Tree} which the lease is taken against 1805 @param notify: L{Smb2} frame containing a LeaseBreakRequest 1806 return a LeaseBreakAcknowledgement with some fields pre-populated 1807 """ 1808 lease_break = notify[0] 1809 smb_req = self.request(obj=tree) 1810 ack_req = smb2.LeaseBreakAcknowledgement(smb_req) 1811 ack_req.lease_key = lease_break.lease_key 1812 ack_req.lease_state = lease_break.new_lease_state 1813 return ack_req

1814

1815 - def oplock_break_acknowledgement(self, fh, notify):

1816 """ 1817 @param fh: Acknowledge break on this L{Open} 1818 @param notify: L{Smb2} frame containing a OplockBreakRequest 1819 return a OplockBreakAcknowledgement with some fields pre-populated 1820 """ 1821 oplock_break = notify[0] 1822 smb_req = self.request(obj=fh) 1823 ack_req = smb2.OplockBreakAcknowledgement(smb_req) 1824 ack_req.file_id = oplock_break.file_id 1825 ack_req.oplock_level = oplock_break.oplock_level 1826 return ack_req

1827

1828 - def frame(self):

1829 return self.connection.frame()

1830

1831 - def request(self, nb=None, obj=None, encrypt_data=None):

1832 smb_req = self.connection.request(nb) 1833 smb_req.session_id = self.session.session_id 1834 1835 if isinstance(obj, Tree): 1836 smb_req.tree_id = obj.tree_id 1837 elif isinstance(obj, Open): 1838 smb_req.tree_id = obj.tree.tree_id 1839 1840 # encryption unspecified, follow session/tree negotiation 1841 if encrypt_data is None: 1842 encrypt_data = self.session.encrypt_data 1843 if isinstance(obj, Tree): 1844 encrypt_data |= obj.encrypt_data 1845 elif isinstance(obj, Open): 1846 encrypt_data |= obj.tree.encrypt_data 1847 1848 # a packet is either encrypted or signed 1849 if encrypt_data and self.session.encryption_context is not None: 1850 transform = crypto.TransformHeader(smb_req.parent) 1851 transform.encryption_context = self.session.encryption_context 1852 transform.session_id = self.session.session_id 1853 elif self.connection.negotiate_response.security_mode & smb2.SMB2_NEGOTIATE_SIGNING_REQUIRED or \ 1854 self.connection.client.security_mode & smb2.SMB2_NEGOTIATE_SIGNING_REQUIRED: 1855 smb_req.flags |= smb2.SMB2_FLAGS_SIGNED 1856 1857 return smb_req

1858

1859 - def let(self, **kwargs):

1860 return self.connection.let(**kwargs)

1861

1862 -class Tree(object):

1863 - def __init__(self, session, path, smb_res):

1864 object.__init__(self) 1865 self.session = session 1866 self.path = path 1867 self.tree_id = smb_res.tree_id 1868 self.tree_connect_response = smb_res[0] 1869 self.encrypt_data = False 1870 if smb_res[0].share_flags & smb2.SMB2_SHAREFLAG_ENCRYPT_DATA: 1871 self.encrypt_data = True 1872 self.session._trees[self.tree_id] = self

1873

1874 -class Open(object):

1875 - def __init__(self, tree, smb_res, create_guid=None, prev=None):

1876 object.__init__(self) 1877 1878 self.create_response = smb_res[0] 1879 1880 self.tree = tree 1881 self.file_id = self.create_response.file_id 1882 self.oplock_level = self.create_response.oplock_level 1883 self.lease = None 1884 self.durable_timeout = None 1885 self.durable_flags = None 1886 self.create_guid = create_guid 1887 1888 if prev is not None: 1889 self.durable_timeout = prev.durable_timeout 1890 self.durable_flags = prev.durable_flags 1891 1892 if self.oplock_level != smb2.SMB2_OPLOCK_LEVEL_NONE: 1893 if self.oplock_level == smb2.SMB2_OPLOCK_LEVEL_LEASE: 1894 lease_res = filter( 1895 lambda c: isinstance(c, smb2.LeaseResponse), 1896 self.create_response)[0] 1897 self.lease = tree.session.client.lease(tree, lease_res) 1898 else: 1899 self.arm_oplock_future() 1900 1901 durable_v2_res = filter( 1902 lambda c: isinstance(c, smb2.DurableHandleV2Response), 1903 self.create_response) 1904 if durable_v2_res != []: 1905 self.durable_timeout = durable_v2_res[0].timeout 1906 self.durable_flags = durable_v2_res[0].flags

1907

1908 - def arm_oplock_future(self):

1909 """ 1910 (Re)arm the oplock future for this open. This function should be called 1911 when an oplock changes level to anything except SMB2_OPLOCK_LEVEL_NONE 1912 """ 1913 self.oplock_future = self.tree.session.client.oplock_break_future( 1914 self.file_id)

1915

1916 - def on_oplock_break(self, cb):

1917 """ 1918 Simple oplock break callback handler. 1919 @param cb: callable taking 1 parameter: the break request oplock level 1920 should return the desired oplock level to break to 1921 """ 1922 def simple_handle_break(op, smb_res, cb_ctx): 1923 """ 1924 note that op is not used in this callback, 1925 since it already closes over self 1926 """ 1927 notify = smb_res[0] 1928 if self.oplock_level != smb2.SMB2_OPLOCK_LEVEL_II: 1929 chan = self.tree.session.first_channel() 1930 ack = chan.oplock_break_acknowledgement(self, smb_res) 1931 ack.oplock_level = cb(notify.oplock_level) 1932 ack_res = chan.connection.transceive(ack.parent.parent)[0][0] 1933 if ack.oplock_level != smb2.SMB2_OPLOCK_LEVEL_NONE: 1934 self.arm_oplock_future() 1935 self.on_oplock_break(cb) 1936 self.oplock_level = ack_res.oplock_level 1937 else: 1938 self.oplock_level = notify.oplock_level

1939 self.on_oplock_break_request(simple_handle_break)

1940

1941 - def on_oplock_break_request(self, cb, cb_ctx=None):

1942 """ 1943 Complex oplock break callback handler. 1944 @param cb: callable taking 3 parameters: 1945 L{Open} 1946 L{Smb2} containing the break request 1947 L{object} arbitrary context 1948 should handle breaking the oplock in some way 1949 callback is also responsible for re-arming the future 1950 and updating the oplock_level (if changed) 1951 """ 1952 def handle_break(f): 1953 smb_res = f.result() 1954 cb(self, smb_res, cb_ctx)

1955 self.oplock_future.then(handle_break) 1956

1957 - def dispose(self):

1958 self.tree = None 1959 if self.lease is not None: 1960 self.lease.dispose() 1961 self.lease = None

1962

1963 -class Lease(object):

1964 - def __init__(self, tree):

1965 self.tree = tree 1966 self.refs = 1 1967 self.future = None

1968

1969 - def update(self, lease_res):

1970 self.lease_key = lease_res.lease_key 1971 self.lease_state = lease_res.lease_state 1972 if self.future is None: 1973 self.arm_future()

1974

1975 - def arm_future(self):

1976 """ 1977 (Re)arm the lease future for this Lease. This function should be called 1978 when a lease changes state to anything other than SMB2_LEASE_NONE 1979 """ 1980 self.future = self.tree.session.client.lease_break_future(self.lease_key)

1981

1982 - def ref(self):

1983 self.refs += 1

1984

1985 - def dispose(self):

1986 self.refs -= 1 1987 if self.refs == 0: 1988 self.tree.session.client.dispose_lease(self)

1989

1990 - def on_break(self, cb):

1991 """ 1992 Simple lease break callback handler. 1993 @param cb: callable taking 1 parameter: the break request lease state 1994 should return the desired lease state to break to 1995 """ 1996 def simple_handle_break(lease, smb_res, cb_ctx): 1997 """ 1998 note that lease is not used in this callback, 1999 since it already closes over self 2000 """ 2001 notify = smb_res[0] 2002 if notify.flags & smb2.SMB2_NOTIFY_BREAK_LEASE_FLAG_ACK_REQUIRED: 2003 chan = self.tree.session.first_channel() 2004 ack = chan.lease_break_acknowledgement(self.tree, smb_res) 2005 ack.lease_state = cb(notify.new_lease_state) 2006 ack_res = chan.connection.transceive(ack.parent.parent)[0][0] 2007 if ack_res.lease_state != smb2.SMB2_LEASE_NONE: 2008 self.arm_future() 2009 self.on_break(cb) 2010 self.lease_state = ack_res.lease_state 2011 else: 2012 self.lease_state = notify.new_lease_state

2013 self.on_break_request(simple_handle_break)

2014

2015 - def on_break_request(self, cb, cb_ctx=None):

2016 """ 2017 Complex lease break callback handler. 2018 @param cb: callable taking 3 parameters: 2019 L{Lease} 2020 L{Smb2} containing the break request 2021 L{object} arbitrary context 2022 should handle breaking the lease in some way 2023 callback is also responsible for re-arming the future 2024 and updating the lease_state (if changed) 2025 """ 2026 def handle_break(f): 2027 smb_res = f.result() 2028 cb(self, smb_res, cb_ctx)

2029 self.future.then(handle_break) 2030

Source Code for Module pike.model