From 7af0827901750b1960bddcdcc49014b73d84cbee Mon Sep 17 00:00:00 2001
From: Arnab Das <arnab.das@emc.com>
Date: Wed, 4 Dec 2019 15:00:44 -0800
Subject: [PATCH] Added pike functions for enum snapshots & sparse

---Review comments addressed
---Cleanup fixtures
--- Renamed to fsctl_request function and bumped version
---Added fixture to catch exception in case open does not have an
appropriate channel to close.
---
 pike/__init__.py |  2 +-
 pike/model.py    | 68 +++++++++++++++++++++++++++++++++++++++++-------
 pike/smb2.py     | 44 +++++++++++++++++++++++++++++++
 3 files changed, 104 insertions(+), 10 deletions(-)

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