From a5bf7aad54b7053417a24ae0918ee42ceed7bf21 Mon Sep 17 00:00:00 2001
From: Joseph Sutton <josephsutton@catalyst.net.nz>
Date: Mon, 13 Sep 2021 20:20:23 +1200
Subject: [PATCH] tests/krb5: Add method to get RODC krbtgt credentials

BUG: https://bugzilla.samba.org/show_bug.cgi?id=14642

Signed-off-by: Joseph Sutton <josephsutton@catalyst.net.nz>
Reviewed-by: Andrew Bartlett <abartlet@samba.org>
Reviewed-by: Douglas Bagnall <douglas.bagnall@catalyst.net.nz>
---
 python/samba/tests/krb5/kdc_base_test.py | 74 ++++++++++++++++++++++++
 1 file changed, 74 insertions(+)

diff --git a/python/samba/tests/krb5/kdc_base_test.py b/python/samba/tests/krb5/kdc_base_test.py
index b2b9d87c3af..fd2e4d08cd3 100644
--- a/python/samba/tests/krb5/kdc_base_test.py
+++ b/python/samba/tests/krb5/kdc_base_test.py
@@ -89,6 +89,7 @@ class KDCBaseTest(RawKerberosTest):
         cls._lp = None
 
         cls._ldb = None
+        cls._rodc_ldb = None
 
         cls._functional_level = None
 
@@ -137,6 +138,30 @@ class KDCBaseTest(RawKerberosTest):
 
         return self._ldb
 
+    def get_rodc_samdb(self):
+        if self._rodc_ldb is None:
+            creds = self.get_admin_creds()
+            lp = self.get_lp()
+
+            session = system_session()
+            type(self)._rodc_ldb = SamDB(url="ldap://%s" % self.host,
+                                         session_info=session,
+                                         credentials=creds,
+                                         lp=lp,
+                                         am_rodc=True)
+
+        return self._rodc_ldb
+
+    def get_server_dn(self, samdb):
+        server = samdb.get_serverName()
+
+        res = samdb.search(base=server,
+                           scope=ldb.SCOPE_BASE,
+                           attrs=['serverReference'])
+        dn = ldb.Dn(samdb, res[0]['serverReference'][0].decode('utf8'))
+
+        return dn
+
     def get_domain_functional_level(self, ldb):
         if self._functional_level is None:
             res = ldb.search(base='',
@@ -478,6 +503,55 @@ class KDCBaseTest(RawKerberosTest):
                                  fallback_creds_fn=create_service_account)
         return c
 
+    def get_rodc_krbtgt_creds(self,
+                              require_keys=True,
+                              require_strongest_key=False):
+        if require_strongest_key:
+            self.assertTrue(require_keys)
+
+        def download_rodc_krbtgt_creds():
+            samdb = self.get_samdb()
+            rodc_samdb = self.get_rodc_samdb()
+
+            rodc_dn = self.get_server_dn(rodc_samdb)
+
+            res = samdb.search(rodc_dn,
+                               scope=ldb.SCOPE_BASE,
+                               attrs=['msDS-KrbTgtLink'])
+            krbtgt_dn = res[0]['msDS-KrbTgtLink'][0]
+
+            res = samdb.search(krbtgt_dn,
+                               scope=ldb.SCOPE_BASE,
+                               attrs=['sAMAccountName',
+                                      'msDS-KeyVersionNumber',
+                                      'msDS-SecondaryKrbTgtNumber'])
+            krbtgt_dn = res[0].dn
+            username = str(res[0]['sAMAccountName'])
+
+            creds = KerberosCredentials()
+            creds.set_domain(self.env_get_var('DOMAIN', 'RODC_KRBTGT'))
+            creds.set_realm(self.env_get_var('REALM', 'RODC_KRBTGT'))
+            creds.set_username(username)
+
+            kvno = int(res[0]['msDS-KeyVersionNumber'][0])
+            krbtgt_number = int(res[0]['msDS-SecondaryKrbTgtNumber'][0])
+
+            rodc_kvno = krbtgt_number << 16 | kvno
+            creds.set_kvno(rodc_kvno)
+            creds.set_dn(krbtgt_dn)
+
+            keys = self.get_keys(samdb, krbtgt_dn)
+            self.creds_set_keys(creds, keys)
+
+            return creds
+
+        c = self._get_krb5_creds(prefix='RODC_KRBTGT',
+                                 allow_missing_password=True,
+                                 allow_missing_keys=not require_keys,
+                                 require_strongest_key=require_strongest_key,
+                                 fallback_creds_fn=download_rodc_krbtgt_creds)
+        return c
+
     def get_krbtgt_creds(self,
                          require_keys=True,
                          require_strongest_key=False):
-- 
2.34.1