tests/krb5: Generify protected users test methods

We can reuse them to test accounts restricted authentication in some
form or another.

Signed-off-by: Joseph Sutton <josephsutton@catalyst.net.nz>
Reviewed-by: Andrew Bartlett <abartlet@samba.org>

3 files changed, 56 insertions, 49 deletions

diff --git a/python/samba/tests/krb5/group_tests.py b/python/samba/tests/krb5/group_tests.py
index c83e4178fa5..9878438e99b 100755
--- a/python/samba/tests/krb5/group_tests.py
+++ b/python/samba/tests/krb5/group_tests.py
@@ -236,7 +236,6 @@ class GroupTests(KDCBaseTest):
         level = netlogon.NetlogonValidationSamInfo
         validation = self._test_samlogon(creds=creds,
                                          logon_type=interactive,
-                                         protected=False,
                                          validation_level=level)
         self.assertIsInstance(validation, netlogon.netr_SamInfo2)
 
@@ -307,7 +306,6 @@ class GroupTests(KDCBaseTest):
         level = netlogon.NetlogonValidationSamInfo2
         validation = self._test_samlogon(creds=creds,
                                          logon_type=interactive,
-                                         protected=False,
                                          validation_level=level)
         self.assertIsInstance(validation, netlogon.netr_SamInfo3)
 
@@ -385,7 +383,6 @@ class GroupTests(KDCBaseTest):
         level = netlogon.NetlogonValidationSamInfo4
         validation = self._test_samlogon(creds=creds,
                                          logon_type=interactive,
-                                         protected=False,
                                          validation_level=level)
         self.assertIsInstance(validation, netlogon.netr_SamInfo6)
 
diff --git a/python/samba/tests/krb5/kdc_base_test.py b/python/samba/tests/krb5/kdc_base_test.py
index f947d45428d..fbc3447bd97 100644
--- a/python/samba/tests/krb5/kdc_base_test.py
+++ b/python/samba/tests/krb5/kdc_base_test.py
@@ -2985,28 +2985,38 @@ class KDCBaseTest(RawKerberosTest):
                                               pac=pac)
 
     # Test credentials by connecting to the DC through LDAP.
-    def _connect(self, creds, expect_error=False):
+    def _connect(self, creds, simple_bind, expect_error=None):
         samdb = self.get_samdb()
+        dn = creds.get_dn()
+
+        if simple_bind:
+            url = f'ldaps://{samdb.host_dns_name()}'
+            creds.set_bind_dn(str(dn))
+        else:
+            url = f'ldap://{samdb.host_dns_name()}'
+            creds.set_bind_dn(None)
         try:
-            ldap = SamDB(url=f'ldap://{samdb.host_dns_name()}',
+            ldap = SamDB(url=url,
                          credentials=creds,
                          lp=self.get_lp())
         except ldb.LdbError as err:
-            self.assertTrue(expect_error, 'got unexpected error')
-            num, _ = err.args
+            self.assertIsNotNone(expect_error, 'got unexpected error')
+            num, estr = err.args
             if num != ldb.ERR_INVALID_CREDENTIALS:
                 raise
 
+            self.assertIn(expect_error, estr)
+
             return
         else:
-            self.assertFalse(expect_error, 'expected to get an error')
+            self.assertIsNone(expect_error, 'expected to get an error')
 
         res = ldap.search('',
                           scope=ldb.SCOPE_BASE,
                           attrs=['tokenGroups'])
         self.assertEqual(1, len(res))
 
-        sid = self.get_objectSid(samdb, creds.get_dn())
+        sid = self.get_objectSid(samdb, dn)
 
         token_groups = res[0].get('tokenGroups', idx=0)
         token_sid = ndr_unpack(security.dom_sid, token_groups)
@@ -3017,7 +3027,7 @@ class KDCBaseTest(RawKerberosTest):
     # user is protected, we should get an ACCOUNT_RESTRICTION error indicating
     # that the password change is not allowed; otherwise we should get a
     # WRONG_PASSWORD error.
-    def _test_samr_change_password(self, creds, protected):
+    def _test_samr_change_password(self, creds, expect_error):
         samdb = self.get_samdb()
         server_name = samdb.host_dns_name()
         conn = samr.samr(f'ncacn_np:{server_name}[krb5,seal,smb2]')
@@ -3033,6 +3043,9 @@ class KDCBaseTest(RawKerberosTest):
         nt_password = samr.CryptPassword()
         nt_verifier = samr.Password()
 
+        if not self.expect_nt_hash:
+            expect_error = ntstatus.NT_STATUS_NTLM_BLOCKED
+
         with self.assertRaises(NTSTATUSError) as err:
             conn.ChangePasswordUser2(server=server,
                                      account=account,
@@ -3043,12 +3056,7 @@ class KDCBaseTest(RawKerberosTest):
                                      lm_verifier=None)
 
         num, _ = err.exception.args
-        if not self.expect_nt_hash:
-            self.assertEqual(ntstatus.NT_STATUS_NTLM_BLOCKED, num)
-        elif protected:
-            self.assertEqual(ntstatus.NT_STATUS_ACCOUNT_RESTRICTION, num)
-        else:
-            self.assertEqual(ntstatus.NT_STATUS_WRONG_PASSWORD, num)
+        self.assertEqual(num, expect_error)
 
         with self.assertRaises(NTSTATUSError) as err:
             conn.ChangePasswordUser3(server=server,
@@ -3061,16 +3069,11 @@ class KDCBaseTest(RawKerberosTest):
                                      password3=None)
 
         num, _ = err.exception.args
-        if not self.expect_nt_hash:
-            self.assertEqual(ntstatus.NT_STATUS_NTLM_BLOCKED, num)
-        elif protected:
-            self.assertEqual(ntstatus.NT_STATUS_ACCOUNT_RESTRICTION, num)
-        else:
-            self.assertEqual(ntstatus.NT_STATUS_WRONG_PASSWORD, num)
+        self.assertEqual(num, expect_error)
 
     # Test SamLogon. Authentication should succeed for non-protected accounts,
     # and fail for protected accounts.
-    def _test_samlogon(self, creds, logon_type, protected,
+    def _test_samlogon(self, creds, logon_type, expect_error=None,
                        validation_level=netlogon.NetlogonValidationSamInfo2):
         samdb = self.get_samdb()
 
@@ -3141,6 +3144,10 @@ class KDCBaseTest(RawKerberosTest):
         netr_flags = 0
 
         validation = None
+
+        if not expect_error and not self.expect_nt_hash:
+            expect_error = ntstatus.NT_STATUS_NTLM_BLOCKED
+
         try:
             (validation, authoritative, flags) = (
                 conn.netr_LogonSamLogonEx(server,
@@ -3150,17 +3157,12 @@ class KDCBaseTest(RawKerberosTest):
                                           validation_level,
                                           netr_flags))
         except NTSTATUSError as err:
-            num, _ = err.args
-            if protected:
-                if num != ntstatus.NT_STATUS_ACCOUNT_RESTRICTION:
-                    raise
-            else:
-                self.assertFalse(self.expect_nt_hash, 'got unexpected error')
-                if num != ntstatus.NT_STATUS_NTLM_BLOCKED:
-                    raise
+            status, _ = err.args
+            self.assertIsNotNone(expect_error,
+                                 f'unexpectedly failed with {status:08X}')
+            self.assertEqual(expect_error, status, 'got wrong status code')
         else:
-            self.assertFalse(protected, 'expected error')
-            self.assertTrue(self.expect_nt_hash, 'expected error')
+            self.assertIsNone(expect_error, 'expected error')
 
             self.assertEqual(1, authoritative)
             self.assertEqual(0, flags)
diff --git a/python/samba/tests/krb5/protected_users_tests.py b/python/samba/tests/krb5/protected_users_tests.py
index c9ed23b5fb9..e29d7c0f6c9 100755
--- a/python/samba/tests/krb5/protected_users_tests.py
+++ b/python/samba/tests/krb5/protected_users_tests.py
@@ -26,7 +26,7 @@ from functools import partial
 
 import ldb
 
-from samba import generate_random_password
+from samba import generate_random_password, ntstatus
 from samba.dcerpc import netlogon, security
 
 import samba.tests.krb5.kcrypto as kcrypto
@@ -51,6 +51,8 @@ import samba.tests.krb5.rfc4120_pyasn1 as krb5_asn1
 global_asn1_print = False
 global_hexdump = False
 
+HRES_SEC_E_LOGON_DENIED = 0x8009030C
+
 
 class ProtectedUsersTests(KDCBaseTest):
     def setUp(self):
@@ -92,7 +94,7 @@ class ProtectedUsersTests(KDCBaseTest):
                                        ntlm=True,
                                        cached=False)
 
-        self._connect(client_creds)
+        self._connect(client_creds, simple_bind=False)
 
     # Test NTLM authentication with a protected account. Authentication should
     # fail, as Protected User accounts cannot use NTLM authentication.
@@ -101,7 +103,8 @@ class ProtectedUsersTests(KDCBaseTest):
                                        ntlm=True,
                                        cached=False)
 
-        self._connect(client_creds, expect_error=True)
+        self._connect(client_creds, simple_bind=False,
+                      expect_error=f'{HRES_SEC_E_LOGON_DENIED:08X}')
 
     # Test that the Protected Users restrictions still apply when the user is a
     # member of a group that is itself a member of Protected Users.
@@ -119,7 +122,8 @@ class ProtectedUsersTests(KDCBaseTest):
                                        ntlm=True,
                                        member_of=group_dn)
 
-        self._connect(client_creds, expect_error=True)
+        self._connect(client_creds, simple_bind=False,
+                      expect_error=f'{HRES_SEC_E_LOGON_DENIED:08X}')
 
     # Test SAMR password changes for unprotected and protected accounts.
     def test_samr_change_password_not_protected(self):
@@ -128,7 +132,9 @@ class ProtectedUsersTests(KDCBaseTest):
         client_creds = self._get_creds(protected=False,
                                        cached=False)
 
-        self._test_samr_change_password(client_creds, protected=False)
+        self._test_samr_change_password(
+            client_creds,
+            expect_error=ntstatus.NT_STATUS_WRONG_PASSWORD)
 
     def test_samr_change_password_protected(self):
         # Use a non-cached account so that it is not locked out for other
@@ -136,39 +142,41 @@ class ProtectedUsersTests(KDCBaseTest):
         client_creds = self._get_creds(protected=True,
                                        cached=False)
 
-        self._test_samr_change_password(client_creds, protected=True)
+        self._test_samr_change_password(
+            client_creds,
+            expect_error=ntstatus.NT_STATUS_ACCOUNT_RESTRICTION)
 
     # Test interactive SamLogon with an unprotected account.
     def test_samlogon_interactive_not_protected(self):
         client_creds = self._get_creds(protected=False,
                                        ntlm=True)
         self._test_samlogon(creds=client_creds,
-                            logon_type=netlogon.NetlogonInteractiveInformation,
-                            protected=False)
+                            logon_type=netlogon.NetlogonInteractiveInformation)
 
     # Test interactive SamLogon with a protected account.
     def test_samlogon_interactive_protected(self):
         client_creds = self._get_creds(protected=True,
                                        ntlm=True)
-        self._test_samlogon(creds=client_creds,
-                            logon_type=netlogon.NetlogonInteractiveInformation,
-                            protected=True)
+        self._test_samlogon(
+            creds=client_creds,
+            logon_type=netlogon.NetlogonInteractiveInformation,
+            expect_error=ntstatus.NT_STATUS_ACCOUNT_RESTRICTION)
 
     # Test network SamLogon with an unprotected account.
     def test_samlogon_network_not_protected(self):
         client_creds = self._get_creds(protected=False,
                                        ntlm=True)
         self._test_samlogon(creds=client_creds,
-                            logon_type=netlogon.NetlogonNetworkInformation,
-                            protected=False)
+                            logon_type=netlogon.NetlogonNetworkInformation)
 
     # Test network SamLogon with a protected account.
     def test_samlogon_network_protected(self):
         client_creds = self._get_creds(protected=True,
                                        ntlm=True)
-        self._test_samlogon(creds=client_creds,
-                            logon_type=netlogon.NetlogonNetworkInformation,
-                            protected=True)
+        self._test_samlogon(
+            creds=client_creds,
+            logon_type=netlogon.NetlogonNetworkInformation,
+            expect_error=ntstatus.NT_STATUS_ACCOUNT_RESTRICTION)
 
     # Test that changing the password of an account in the Protected Users
     # group still generates an NT hash.