From 12a1fabd121b2d54d94ed971c3af0c6c3b3d59c7 Mon Sep 17 00:00:00 2001 From: Joseph Sutton Date: Tue, 7 Mar 2023 15:28:21 +1300 Subject: [PATCH] tests/krb5: Cache drsuapi connection We call get_keys() a lot, and it's more efficient if we aren't creating a new connection for every new account we create. To allow us to maintain a single cached connection, remove the samdb parameter from get_keys() and get_secrets(). No-one was using it anyway. Signed-off-by: Joseph Sutton Reviewed-by: Andrew Bartlett --- python/samba/tests/krb5/kdc_base_test.py | 41 +++++++++++-------- python/samba/tests/krb5/nt_hash_tests.py | 1 - .../samba/tests/krb5/protected_users_tests.py | 2 +- python/samba/tests/krb5/spn_tests.py | 2 +- 4 files changed, 26 insertions(+), 20 deletions(-) diff --git a/python/samba/tests/krb5/kdc_base_test.py b/python/samba/tests/krb5/kdc_base_test.py index 6500a188972..18ee8f8bd98 100644 --- a/python/samba/tests/krb5/kdc_base_test.py +++ b/python/samba/tests/krb5/kdc_base_test.py @@ -151,6 +151,8 @@ class KDCBaseTest(RawKerberosTest): cls._ldb = None cls._rodc_ldb = None + cls._drsuapi_connection = None + cls._functional_level = None # An identifier to ensure created accounts have unique names. Windows @@ -281,6 +283,18 @@ class KDCBaseTest(RawKerberosTest): return self._rodc_ldb + def get_drsuapi_connection(self): + if self._drsuapi_connection is None: + admin_creds = self.get_admin_creds() + samdb = self.get_samdb() + dns_hostname = samdb.host_dns_name() + type(self)._drsuapi_connection = drsuapi_connect(dns_hostname, + self.get_lp(), + admin_creds, + ip=self.dc_host) + + return self._drsuapi_connection + def get_server_dn(self, samdb): server = samdb.get_serverName() @@ -686,7 +700,6 @@ class KDCBaseTest(RawKerberosTest): rodc_ctx = self.get_mock_rodc_ctx() self.get_secrets( - samdb, dn, destination_dsa_guid=rodc_ctx.ntds_guid, source_dsa_invocation_id=misc.GUID(samdb.invocation_id)) @@ -712,16 +725,10 @@ class KDCBaseTest(RawKerberosTest): else: self.assertNotIn(str(dn), revealed_dns) - def get_secrets(self, samdb, dn, + def get_secrets(self, dn, destination_dsa_guid, source_dsa_invocation_id): - admin_creds = self.get_admin_creds() - - dns_hostname = samdb.host_dns_name() - (bind, handle, _) = drsuapi_connect(dns_hostname, - self.get_lp(), - admin_creds, - ip=self.dc_host) + bind, handle, _ = self.get_drsuapi_connection() req = drsuapi.DsGetNCChangesRequest8() @@ -773,11 +780,11 @@ class KDCBaseTest(RawKerberosTest): return bind, identifier, attributes - def get_keys(self, samdb, dn, expected_etypes=None): + def get_keys(self, dn, expected_etypes=None): admin_creds = self.get_admin_creds() + samdb = self.get_samdb() bind, identifier, attributes = self.get_secrets( - samdb, str(dn), destination_dsa_guid=misc.GUID(samdb.get_ntds_GUID()), source_dsa_invocation_id=misc.GUID()) @@ -1444,7 +1451,7 @@ class KDCBaseTest(RawKerberosTest): expected_etypes = None if force_nt4_hash: expected_etypes = {kcrypto.Enctype.RC4} - keys = self.get_keys(samdb, dn, expected_etypes=expected_etypes) + keys = self.get_keys(dn, expected_etypes=expected_etypes) self.creds_set_keys(creds, keys) # Handle secret replication to the RODC. @@ -1628,7 +1635,7 @@ class KDCBaseTest(RawKerberosTest): creds.set_kvno(rodc_kvno) creds.set_dn(krbtgt_dn) - keys = self.get_keys(samdb, krbtgt_dn) + keys = self.get_keys(krbtgt_dn) self.creds_set_keys(creds, keys) # The RODC krbtgt account should support the default enctypes, @@ -1681,7 +1688,7 @@ class KDCBaseTest(RawKerberosTest): creds.set_kvno(rodc_kvno) creds.set_dn(dn) - keys = self.get_keys(samdb, dn) + keys = self.get_keys(dn) self.creds_set_keys(creds, keys) extra_bits = (security.KERB_ENCTYPE_AES128_CTS_HMAC_SHA1_96 | @@ -1729,7 +1736,7 @@ class KDCBaseTest(RawKerberosTest): creds.set_kvno(kvno) creds.set_dn(dn) - keys = self.get_keys(samdb, dn) + keys = self.get_keys(dn) self.creds_set_keys(creds, keys) # The krbtgt account should support the default enctypes, although @@ -1780,7 +1787,7 @@ class KDCBaseTest(RawKerberosTest): creds.set_workstation(username[:-1]) creds.set_dn(dn) - keys = self.get_keys(samdb, dn) + keys = self.get_keys(dn) self.creds_set_keys(creds, keys) extra_bits = (security.KERB_ENCTYPE_AES128_CTS_HMAC_SHA1_96 | @@ -1827,7 +1834,7 @@ class KDCBaseTest(RawKerberosTest): creds.set_kvno(kvno) creds.set_dn(dn) - keys = self.get_keys(samdb, dn) + keys = self.get_keys(dn) self.creds_set_keys(creds, keys) extra_bits = (security.KERB_ENCTYPE_AES128_CTS_HMAC_SHA1_96 | diff --git a/python/samba/tests/krb5/nt_hash_tests.py b/python/samba/tests/krb5/nt_hash_tests.py index f2cd14887f8..82d9c09eb86 100755 --- a/python/samba/tests/krb5/nt_hash_tests.py +++ b/python/samba/tests/krb5/nt_hash_tests.py @@ -47,7 +47,6 @@ class NtHashTests(KDCBaseTest): admin_creds = self.get_admin_creds() bind, identifier, attributes = self.get_secrets( - samdb, dn, destination_dsa_guid=misc.GUID(samdb.get_ntds_GUID()), source_dsa_invocation_id=misc.GUID()) diff --git a/python/samba/tests/krb5/protected_users_tests.py b/python/samba/tests/krb5/protected_users_tests.py index af6b6b57bf3..5ca5de0321e 100755 --- a/python/samba/tests/krb5/protected_users_tests.py +++ b/python/samba/tests/krb5/protected_users_tests.py @@ -291,7 +291,7 @@ class ProtectedUsersTests(KDCBaseTest): client_creds.set_password(new_password) - self.get_keys(samdb, client_dn, + self.get_keys(client_dn, expected_etypes={kcrypto.Enctype.AES256, kcrypto.Enctype.AES128, kcrypto.Enctype.RC4}) diff --git a/python/samba/tests/krb5/spn_tests.py b/python/samba/tests/krb5/spn_tests.py index f4f20bea4f2..42e7b3428d6 100755 --- a/python/samba/tests/krb5/spn_tests.py +++ b/python/samba/tests/krb5/spn_tests.py @@ -199,7 +199,7 @@ class SpnTests(KDCBaseTest): kvno = int(res[0].get('msDS-KeyVersionNumber', idx=0)) creds.set_kvno(kvno) - keys = self.get_keys(samdb, rodc_dn) + keys = self.get_keys(rodc_dn) self.creds_set_keys(creds, keys) return creds