diff --git a/src/pytds/__init__.py b/src/pytds/__init__.py index e019cd2..577ef83 100644 --- a/src/pytds/__init__.py +++ b/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 diff --git a/src/pytds/tds.py b/src/pytds/tds.py index 02c50e3..e6c51a1 100644 --- a/src/pytds/tds.py +++ b/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 = { diff --git a/src/pytds/tds_base.py b/src/pytds/tds_base.py index 59480ca..345e146 100644 --- a/src/pytds/tds_base.py +++ b/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 #