Merge pull request #87 from isi-adas/BR_ISILONQE_977

Br isilonqe 977

pike/model.py

@@ -361,6 +361,28 @@ def unregister_callback(self, event, cb):
             return
         self.callbacks[ev].remove(cb)
 
+    def restrict_dialects(self, min_dialect=None, max_dialect=None):
+        """
+        Update the list of SMB2 dialects this Client will advertise support for.
+
+        This call will only remove dialects from the list.
+
+        Dialects should be from pike.smb2.Dialect enum, but technically may be any
+        int or float.
+
+        If either of min_dialect or max_dialect is None, the minimum or maximum dialect
+        supported will be used.
+
+        @param min_dialect: The minimum dialect to support (inclusive)
+        @param max_dialect: The maximum dialect to support (inclusive)
+        """
+        if min_dialect is None:
+            min_dialect = min(smb2.Dialect.values())
+        if max_dialect is None:
+            max_dialect = max(smb2.Dialect.values())
+        self.dialects = [d for d in self.dialects
+                         if min_dialect <= d <= max_dialect]
+
     def connect(self, server, port=445):
         """
         Create a connection.

pike/test/__init__.py

@@ -160,8 +160,7 @@ def default_client():
     client = model.Client()
     min_dialect = Options.min_dialect()
     max_dialect = Options.max_dialect()
-    client.dialects = [d for d in client.dialects
-                       if min_dialect <= d <= max_dialect]
+    client.restrict_dialects(min_dialect, max_dialect)
     if Options.signing():
         client.security_mode = (smb2.SMB2_NEGOTIATE_SIGNING_ENABLED |
                                 smb2.SMB2_NEGOTIATE_SIGNING_REQUIRED)
@@ -408,6 +407,12 @@ def error(self, *args, **kwargs):
     def critical(self, *args, **kwargs):
         self.logger.critical(*args, **kwargs)
 
+    def set_client_dialect(self, min_dialect=None, max_dialect=None,
+                           client=None):
+        if client is None:
+            client = self.default_client
+        client.restrict_dialects(min_dialect, max_dialect)
+
     def tree_connect(self, client=None, resume=None):
         tc = TreeConnect(
             client=client or self.default_client,

pike/test/internal/test_client.py

@@ -0,0 +1,34 @@
+import pytest
+
+from pike.model import Client
+from pike.smb2 import Dialect
+
+
+@pytest.mark.parametrize(
+    "min_dialect,max_dialect",
+    (
+        pytest.param(None, None, id="None/None"),
+        pytest.param(Dialect.DIALECT_SMB2_1, None, id="lower bound"),
+        pytest.param(None, Dialect.DIALECT_SMB2_1, id="upper bound"),
+        pytest.param(Dialect.DIALECT_SMB3_1_1, Dialect.DIALECT_SMB2_002, id="inverted"),
+        pytest.param(Dialect.DIALECT_SMB3_0, float("inf"), id="upper bound inf"),
+        pytest.param(Dialect.DIALECT_SMB3_0, 0xFFFFFFFF, id="upper bound int"),
+        pytest.param(float("-inf"), Dialect.DIALECT_SMB2_002, id="lower bound -inf"),
+        pytest.param(0, Dialect.DIALECT_SMB3_0_2, id="lower bound zero"),
+    ),
+)
+def test_restrict_dialect(min_dialect, max_dialect):
+    c = Client()
+    c.restrict_dialects(min_dialect, max_dialect)
+    print(
+        "restrict_dialects(min_dialect={}, max_dialect={}) => {}".format(
+            min_dialect, max_dialect, c.dialects
+        )
+    )
+    for d in c.dialects:
+        if min_dialect is None:
+            min_dialect = min(Dialect.values())
+        if max_dialect is None:
+            max_dialect = max(Dialect.values())
+        assert d >= min_dialect
+        assert d <= max_dialect

pike/test/ioctl.py

@@ -29,7 +29,7 @@
 #
 # Abstract:
 #
-#        Test IOCTLs.  Currently, only FSCTL_VALIDATE_NEGOTIATE_INFO
+#        Test IOCTLs
 #
 # Authors: Ki Anderson ([email protected])
 #          Paul Martin ([email protected])
@@ -41,29 +41,28 @@
 import pike.test as test
 import pike.ntstatus
 
+
 class ValidateNegotiateInfo(test.PikeTest):
     def __init__(self, *args, **kwds):
         super(ValidateNegotiateInfo, self).__init__(*args, **kwds)
-        self.default_client.dialects = [
-                smb2.DIALECT_SMB3_0,
-                smb2.DIALECT_SMB3_0_2]
+        self.set_client_dialect(smb2.DIALECT_SMB3_0, smb2.DIALECT_SMB3_0_2)
 
     # VALIDATE_NEGOTIATE_INFO fsctl succeeds for SMB3
     @test.RequireDialect(smb2.DIALECT_SMB3_0, smb2.DIALECT_SMB3_0_2)
     def test_validate_negotiate_smb3(self):
+        self.set_client_dialect(smb2.DIALECT_SMB3_0, smb2.DIALECT_SMB3_0_2)
         chan, tree = self.tree_connect()
         chan.validate_negotiate_info(tree)
 
 
 class QueryNetworkInterfaceInfo(test.PikeTest):
     def __init__(self, *args, **kwds):
         super(QueryNetworkInterfaceInfo, self).__init__(*args, **kwds)
-        self.default_client.dialects = [
-                smb2.DIALECT_SMB3_0,
-                smb2.DIALECT_SMB3_0_2]
+        self.set_client_dialect(smb2.DIALECT_SMB3_0, smb2.DIALECT_SMB3_1_1)
 
-    @test.RequireDialect(smb2.DIALECT_SMB3_0, smb2.DIALECT_SMB3_0_2)
+    @test.RequireDialect(smb2.DIALECT_SMB3_0, smb2.DIALECT_SMB3_1_1)
     def test_query_interface_info_smb3(self):
+        self.set_client_dialect(smb2.DIALECT_SMB3_0, smb2.DIALECT_SMB3_1_1)
         chan, tree = self.tree_connect()
         try:
             chan.query_network_interface_info(tree)