tests/krb5: Add helper function to modify ticket flags

Signed-off-by: Joseph Sutton <josephsutton@catalyst.net.nz>
Reviewed-by: Stefan Metzmacher <metze@samba.org>
(cherry picked from commit ded5115f73dff5b8b2f3212988e03f9dbe0c2aa3)

3 files changed, 19 insertions, 30 deletions

diff --git a/python/samba/tests/krb5/kdc_base_test.py b/python/samba/tests/krb5/kdc_base_test.py
index 9506048ee2a..58b87eab25b 100644
--- a/python/samba/tests/krb5/kdc_base_test.py
+++ b/python/samba/tests/krb5/kdc_base_test.py
@@ -1602,6 +1602,20 @@ class KDCBaseTest(RawKerberosTest):
             enc_part, asn1Spec=krb5_asn1.EncTicketPart())
         return enc_ticket_part
 
+    def modify_ticket_flag(self, enc_part, flag, value):
+        self.assertIsInstance(value, bool)
+
+        flag = krb5_asn1.TicketFlags(flag)
+        pos = len(tuple(flag)) - 1
+
+        flags = enc_part['flags']
+        self.assertLessEqual(pos, len(flags))
+
+        new_flags = flags[:pos] + str(int(value)) + flags[pos + 1:]
+        enc_part['flags'] = new_flags
+
+        return enc_part
+
     def get_objectSid(self, samdb, dn):
         ''' Get the objectSID for a DN
             Note: performs an Ldb query.
diff --git a/python/samba/tests/krb5/kdc_tgs_tests.py b/python/samba/tests/krb5/kdc_tgs_tests.py
index 2923d53772a..8cd27dec2aa 100755
--- a/python/samba/tests/krb5/kdc_tgs_tests.py
+++ b/python/samba/tests/krb5/kdc_tgs_tests.py
@@ -2177,14 +2177,7 @@ class KdcTgsTests(KDCBaseTest):
 
     def _modify_renewable(self, enc_part):
         # Set the renewable flag.
-        renewable_flag = krb5_asn1.TicketFlags('renewable')
-        pos = len(tuple(renewable_flag)) - 1
-
-        flags = enc_part['flags']
-        self.assertLessEqual(pos, len(flags))
-
-        new_flags = flags[:pos] + '1' + flags[pos + 1:]
-        enc_part['flags'] = new_flags
+        enc_part = self.modify_ticket_flag(enc_part, 'renewable', value=True)
 
         # Set the renew-till time to be in the future.
         renew_till = self.get_KerberosTime(offset=100 * 60 * 60)
@@ -2194,14 +2187,7 @@ class KdcTgsTests(KDCBaseTest):
 
     def _modify_invalid(self, enc_part):
         # Set the invalid flag.
-        invalid_flag = krb5_asn1.TicketFlags('invalid')
-        pos = len(tuple(invalid_flag)) - 1
-
-        flags = enc_part['flags']
-        self.assertLessEqual(pos, len(flags))
-
-        new_flags = flags[:pos] + '1' + flags[pos + 1:]
-        enc_part['flags'] = new_flags
+        enc_part = self.modify_ticket_flag(enc_part, 'invalid', value=True)
 
         # Set the ticket start time to be in the past.
         past_time = self.get_KerberosTime(offset=-100 * 60 * 60)
diff --git a/python/samba/tests/krb5/s4u_tests.py b/python/samba/tests/krb5/s4u_tests.py
index 6ec9af11423..49dd89cd764 100755
--- a/python/samba/tests/krb5/s4u_tests.py
+++ b/python/samba/tests/krb5/s4u_tests.py
@@ -1336,20 +1336,9 @@ class S4UKerberosTests(KDCBaseTest):
                                     modify_pac_fn=modify_pac_fn)
 
     def set_ticket_forwardable(self, ticket, flag, update_pac_checksums=True):
-        flag = '1' if flag else '0'
-
-        def modify_fn(enc_part):
-            # Reset the forwardable flag
-            forwardable_pos = (len(tuple(krb5_asn1.TicketFlags('forwardable')))
-                               - 1)
-
-            flags = enc_part['flags']
-            self.assertLessEqual(forwardable_pos, len(flags))
-            enc_part['flags'] = (flags[:forwardable_pos] +
-                                 flag +
-                                 flags[forwardable_pos+1:])
-
-            return enc_part
+        modify_fn = functools.partial(self.modify_ticket_flag,
+                                      flag='forwardable',
+                                      value=flag)
 
         if update_pac_checksums:
             checksum_keys = self.get_krbtgt_checksum_key()