Merge pull request #118 from corenec/b.pike.p128466

Add support for V2 SMB leases so that epoch checking can be done.

pike/model.py

@@ -1466,6 +1466,7 @@ def create_request(
         oplock_level=smb2.SMB2_OPLOCK_LEVEL_NONE,
         lease_key=None,
         lease_state=None,
+        lease_flags_v2=None,
         durable=False,
         persistent=False,
         create_guid=None,
@@ -1498,7 +1499,7 @@ def create_request(
                 lease_key = crypto.random_bytes(16)
             if lease_state is None:
                 lease_state = smb2.SMB2_LEASE_RWH
-            lease_req = smb2.LeaseRequest(create_req)
+            lease_req = smb2.LeaseRequest(create_req, lease_flags_v2)
             lease_req.lease_key = lease_key
             lease_req.lease_state = lease_state
 
@@ -1584,6 +1585,7 @@ def create(
         oplock_level=smb2.SMB2_OPLOCK_LEVEL_NONE,
         lease_key=None,
         lease_state=None,
+        lease_flags_v2=None,
         durable=False,
         persistent=False,
         create_guid=None,
@@ -1605,6 +1607,7 @@ def create(
                 oplock_level,
                 lease_key,
                 lease_state,
+                lease_flags_v2,
                 durable,
                 persistent,
                 create_guid,
@@ -2332,6 +2335,7 @@ def __init__(self, tree):
     def update(self, lease_res):
         self.lease_key = lease_res.lease_key
         self.lease_state = lease_res.lease_state
+        self.lease_epoch = lease_res.epoch
         if self.future is None:
             self.arm_future()
 
@@ -2363,6 +2367,7 @@ def simple_handle_break(lease, smb_res, cb_ctx):
             since it already closes over self
             """
             notify = smb_res[0]
+            self.lease_epoch = notify.new_epoch
             if notify.flags & smb2.SMB2_NOTIFY_BREAK_LEASE_FLAG_ACK_REQUIRED:
                 chan = self.tree.session.first_channel()
                 ack = chan.lease_break_acknowledgement(self.tree, smb_res)

pike/smb2.py

@@ -1353,7 +1353,7 @@ class LeaseState(core.FlagEnum):
 class LeaseFlags(core.FlagEnum):
     SMB2_LEASE_FLAG_NONE = 0x00
     SMB2_LEASE_FLAG_BREAK_IN_PROGRESS = 0x02
-
+    SMB2_LEASE_FLAG_PARENT_LEASE_KEY_SET = 0x04
 
 LeaseFlags.import_items(globals())
 
@@ -1829,14 +1829,19 @@ class LeaseRequest(CreateRequestContext):
     # This class handles V2 requests as well.  Set
     # the lease_flags field to a non-None value
     # to enable the extended fields
-    def __init__(self, parent):
+    def __init__(self, parent, lease_flags_v2=None):
         CreateRequestContext.__init__(self, parent)
         self.lease_key = array.array("B", [0] * 16)
         self.lease_state = 0
         # V2 fields
-        self.lease_flags = None
-        self.parent_lease_key = None
-        self.epoch = None
+        if lease_flags_v2 is not None:
+            self.lease_flags = lease_flags_v2
+            self.parent_lease_key = array.array("B", [0] * 16)
+            self.epoch = 0
+        else:
+            self.lease_flags = None
+            self.parent_lease_key = None
+            self.epoch = None 
 
     def _encode(self, cur):
         cur.encode_bytes(self.lease_key)