Federated Auth

src/pytds/__init__.py

@@ -1158,6 +1158,7 @@ def connect(dsn=None, database=None, user=None, password=None, timeout=None,
             enc_login_only=False, disable_connect_retry=False,
             pooling=False,
             use_sso=False,
+            access_token = None
             ):
     """
     Opens connection to the database
@@ -1215,6 +1216,8 @@ def connect(dsn=None, database=None, user=None, password=None, timeout=None,
     :type enc_login_only: bool
     :keyword use_sso: Enables SSO login, e.g. Kerberos using SSPI on Windows and kerberos package on other platforms.
              Cannot be used together with auth parameter.
+    :keyword access_token: Use Azure Active Directory Authentication / Federated Authentication
+    :type access_token: str
     :returns: An instance of :class:`Connection`
     """
     if use_sso and auth:
@@ -1235,6 +1238,7 @@ def connect(dsn=None, database=None, user=None, password=None, timeout=None,
     login.bulk_copy = False
     login.client_lcid = lcid.LANGID_ENGLISH_US
     login.use_mars = use_mars
+    login.access_token = access_token
     login.pid = os.getpid()
     login.change_password = ''
     login.client_id = uuid.getnode()  # client mac address

src/pytds/tds.py

@@ -1229,6 +1229,21 @@ def send_prelogin(self, login):
         instance_name = instance_name.encode('ascii')
         if len(instance_name) > 65490:
             raise ValueError('Instance name is too long')
+        if tds_base.IS_TDS74_PLUS(self):
+            start_pos = 26
+            buf = struct.pack(b'>BHHBHHBHHBHHBHHHB',
+                PreLoginToken.VERSION, start_pos, 6,
+                # encryption
+                PreLoginToken.ENCRYPTION, start_pos + 6, 1,
+                # instance
+                PreLoginToken.INSTOPT, start_pos + 6 + 1, len(instance_name) + 1,
+                # thread id
+                PreLoginToken.THREADID, start_pos + 6 + 1 + len(instance_name) + 1, 4,
+                # MARS enabled
+                PreLoginToken.MARS, start_pos + 6 + 1 + len(instance_name) + 1 + 4, 1,
+                PreLoginToken.B_FEDAUTHREQUIRED, start_pos + start_pos + 6 + 1 + len(instance_name) + 1 + 4 + 1,1,
+                # end
+                PreLoginToken.TERMINATOR)
         if tds_base.IS_TDS72_PLUS(self):
             start_pos = 26
             buf = struct.pack(
@@ -1281,6 +1296,9 @@ def send_prelogin(self, login):
             # MARS (1 enabled)
             w.put_byte(1 if login.use_mars else 0)
             attribs['mars'] = login.use_mars
+        if tds_base.IS_TDS74_PLUS(self):
+            w.put_byte(1 if login.access_token else 0)
+            attribs['fedautch'] = bool(login.access_token)
         logger.info('Sending PRELOGIN %s', ' '.join('%s=%s' % (n, v) for n, v in attribs.items()))
 
         w.flush()
@@ -1323,6 +1341,9 @@ def parse_prelogin(self, octets, login):
             elif type_id == PreLoginToken.INSTOPT:
                 # ignore instance name mismatch
                 pass
+            elif type_id == PreLoginToken.FEDAUTHREQUIRED:
+                if not login.access_token:
+                    raise tds_base.Error('Server requires Federated Auth but was not provided')
             i += 5
         logger.info("Got PRELOGIN response crypt=%x mars=%d",
                     crypt_flag, self.conn._mars_enabled)
@@ -1420,7 +1441,7 @@ def tds7_send_login(self, login):
         w.put_smallint(current_pos)
         w.put_smallint(len(client_host_name))
         current_pos += len(client_host_name) * 2
-        if self.authentication:
+        if self.authentication or self.login.access_token:
             w.put_smallint(0)
             w.put_smallint(0)
             w.put_smallint(0)
@@ -1472,7 +1493,7 @@ def tds7_send_login(self, login):
             # sspi long
             w.put_int(0)
         w.write_ucs2(client_host_name)
-        if not self.authentication:
+        if not self.authentication and not self.login.access_token:
             w.write_ucs2(user_name)
             w.write(tds7_crypt_pass(login.password))
         w.write_ucs2(login.app_name)
@@ -1484,6 +1505,9 @@ def tds7_send_login(self, login):
             w.write(auth_packet)
         w.write_ucs2(login.attach_db_file)
         w.write_ucs2(login.change_password)
+        if login.access_token:
+            w.put_byte(tds_base.TDS_LOGIN_FEATURE_FEDAUTH)
+
         w.flush()
 
     _SERVER_TO_CLIENT_MAPPING = {

src/pytds/tds_base.py

@@ -17,6 +17,7 @@
 IS_TDS71_PLUS = lambda x: x.tds_version >= TDS71
 IS_TDS72_PLUS = lambda x: x.tds_version >= TDS72
 IS_TDS73_PLUS = lambda x: x.tds_version >= TDS73A
+IS_TDS74_PLUS = lambda x: x.tds_version >= TDS74
 
 
 # https://msdn.microsoft.com/en-us/library/dd304214.aspx
@@ -220,6 +221,17 @@ class PacketType:
 TDS_FOLEDB = 0x10
 TDS_FREADONLY_INTENT = 0x20
 
+# TDS Login Features
+# as per https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-tds/773a62b6-ee89-4c02-9e5e-344882630aac
+TDS_LOGIN_FEATURE_SESSIONRECOVERY = 0x01
+TDS_LOGIN_FEATURE_FEDAUTH = 0x02
+TDS_LOGIN_FEATURE_COLUMNENCRYPTION = 0x04
+TDS_LOGIN_FEATURE_GLOBALTRANSACTIONS = 0x05
+TDS_LOGIN_FEATURE_AZURESQLSUPPORT = 0x08
+TDS_LOGIN_FEATURE_DATACLASSIFICATION = 0x09
+TDS_LOGIN_FEATURE_UTF8_SUPPORT = 0x0A
+TDS_LOGIN_FEATURE_AZURESQLDNSCACHING = 0x0B
+
 #
 # Sybase only types
 #