diff --git a/pike/__init__.py b/pike/__init__.py index c4d97c04..8fc3da72 100644 --- a/pike/__init__.py +++ b/pike/__init__.py @@ -13,5 +13,5 @@ 'test', 'transport', ] -__version_info__ = (0, 2, 22) +__version_info__ = (0, 2, 23) __version__ = "{0}.{1}.{2}".format(*__version_info__) diff --git a/pike/model.py b/pike/model.py index c8e4894a..4a01c7da 100644 --- a/pike/model.py +++ b/pike/model.py @@ -1870,21 +1870,58 @@ def get_symlink(self, file): return self.connection.transceive( self.get_symlink_request(file).parent.parent)[0] - def enumerate_snapshots_request(self, fh, max_output_response=16384): - smb_req = self.request(obj=fh.tree) + def fsctl_request(self, smb_req, fh, buflen=16384): ioctl_req = smb2.IoctlRequest(smb_req) - ioctl_req.max_output_response = max_output_response + ioctl_req.max_output_response = buflen + ioctl_req.flags = smb2.SMB2_0_IOCTL_IS_FSCTL ioctl_req.file_id = fh.file_id - ioctl_req.flags |= smb2.SMB2_0_IOCTL_IS_FSCTL - enum_req = smb2.EnumerateSnapshotsRequest(ioctl_req) + return ioctl_req + + def enumerate_snapshots_request(self, fh, + snap_request=smb2.EnumerateSnapshotsRequest, + max_output_response=16384): + smb_req = self.request(obj=fh.tree) + fsctl_req = self.fsctl_request(smb_req, fh, + buflen=max_output_response) + enum_req = snap_request(fsctl_req) return enum_req - def enumerate_snapshots(self, fh): + def enumerate_snapshots(self, fh, snap_request=smb2.EnumerateSnapshotsRequest): return self.connection.transceive( - self.enumerate_snapshots_request(fh).parent.parent.parent)[0] + self.enumerate_snapshots_request(fh, + snap_request).parent.parent.parent)[0] + + def enumerate_snapshots_list(self, fh, + snap_request=smb2.EnumerateSnapshotsRequest): + return self.enumerate_snapshots(fh, snap_request)[0][0].snapshots + + def zero_data(self, tree, src_offsets, dst_offsets, src_filename): + """ Send a FSCTL_SET_ZERO_DATA ioctl request """ + fh_src = self.create(tree, src_filename, + access=smb2.FILE_READ_DATA | smb2.FILE_WRITE_DATA, + share=(smb2.FILE_SHARE_READ | + smb2.FILE_SHARE_WRITE | + smb2.FILE_SHARE_DELETE), + disposition=smb2.FILE_OPEN_IF).result() + smb_req = self.request(obj=tree) + fsctl_req = self.fsctl_request(smb_req, fh_src) + smb2.SetSparseRequest(fsctl_req) + try: + results = self.connection.transceive(smb_req.parent) + except Exception: + self.close(fh_src) + raise - def enumerate_snapshots_list(self, fh): - return self.enumerate_snapshots(fh)[0][0].snapshots + smb_req = self.request(obj=tree) + fsctl_req = self.fsctl_request(smb_req, fh_src) + zerodata_req = smb2.SetZeroDataRequest(fsctl_req) + zerodata_req.file_offset = src_offsets + zerodata_req.beyond_final_zero = dst_offsets + try: + results = self.connection.transceive(smb_req.parent) + return results + finally: + self.close(fh_src) def lease_break_acknowledgement(self, tree, notify): """ @@ -2065,6 +2102,19 @@ def dispose(self): self.lease.dispose() self.lease = None + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + try: + chan = self.tree.session.first_channel() + chan.close(self) + except StopIteration: + # If the underlying connection for the channel is closed explicitly + # open will not able to find an appropriate channel, to send close. + pass + + class Lease(object): def __init__(self, tree): self.tree = tree diff --git a/pike/smb2.py b/pike/smb2.py index 6b6c245a..c1eaddce 100644 --- a/pike/smb2.py +++ b/pike/smb2.py @@ -3843,3 +3843,47 @@ def _decode(self, cur): # TODO: handle the case where the buffer is too small self.snapshots = snapshot_buffer.strip("\0").split("\0") +class SetZeroDataRequest(IoctlInput): + ioctl_ctl_code = IoctlCode.FSCTL_SET_ZERO_DATA + + def __init__(self, parent): + IoctlInput.__init__(self, parent) + self.file_offset = 0 + self.beyond_final_zero = 0 + parent.ioctl_input = self + + def _encode(self, cur): + cur.encode_uint64le(self.file_offset) + cur.encode_uint64le(self.beyond_final_zero) + + +class SetZeroDataResponse(IoctlOutput): + ioctl_ctl_code = IoctlCode.FSCTL_SET_ZERO_DATA + + def __init__(self, parent): + IoctlOutput.__init__(self, parent) + parent.ioctl_output = self + + def _decode(self, cur): + pass + +class SetSparseRequest(IoctlInput): + ioctl_ctl_code = IoctlCode.FSCTL_SET_SPARSE + + def __init__(self, parent): + IoctlInput.__init__(self, parent) + parent.ioctl_input = self + + def _encode(self, cur): + pass + + +class SetSparseResponse(IoctlOutput): + ioctl_ctl_code = IoctlCode.FSCTL_SET_SPARSE + + def __init__(self, parent): + IoctlOutput.__init__(self, parent) + parent.ioctl_output = self + + def _decode(self, cur): + pass