-
Notifications
You must be signed in to change notification settings - Fork 27
FEAT: Adding authentication module and adding new auth types #135
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Merged
Changes from all commits
Commits
Show all changes
15 commits
Select commit
Hold shift + click to select a range
d2f2397
FEAT: Adding authentication module and adding new auth types
jahnvi480 3320013
Resolving comments
jahnvi480 7eaeebb
Resolving comments
jahnvi480 d8b5bfd
Adding comments
jahnvi480 9262f73
Altering testcases
jahnvi480 be1fba4
removed ENABLE_LOGGING
bewithgaurav 1418dc0
Merge branch 'main' of https://github.com/microsoft/mssql-python into…
jahnvi480 6827dbc
fixed conflicts again
bewithgaurav f093eaf
restored others to this branch - reverting back
bewithgaurav d260e52
restoring constants and setup
bewithgaurav fdadbca
Adding setup dependencies
jahnvi480 0b602c8
Merge branch 'jahnvi/new_auth_module' of https://github.com/microsoft…
jahnvi480 992aa6d
Resolving conflicts
jahnvi480 2c61a41
Resolving comments
jahnvi480 5f98adb
Merge branch 'main' into jahnvi/new_auth_module
jahnvi480 File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,163 @@ | ||
| """ | ||
| Copyright (c) Microsoft Corporation. | ||
| Licensed under the MIT license. | ||
| This module handles authentication for the mssql_python package. | ||
| """ | ||
|
|
||
| import platform | ||
| import struct | ||
| from typing import Tuple, Dict, Optional, Union | ||
| from mssql_python.logging_config import get_logger, ENABLE_LOGGING | ||
| from mssql_python.constants import AuthType | ||
|
|
||
| logger = get_logger() | ||
|
|
||
| class AADAuth: | ||
| """Handles Azure Active Directory authentication""" | ||
|
|
||
| @staticmethod | ||
| def get_token_struct(token: str) -> bytes: | ||
| """Convert token to SQL Server compatible format""" | ||
| token_bytes = token.encode("UTF-16-LE") | ||
| return struct.pack(f"<I{len(token_bytes)}s", len(token_bytes), token_bytes) | ||
|
|
||
| @staticmethod | ||
| def get_token(auth_type: str) -> bytes: | ||
| """Get token using the specified authentication type""" | ||
| from azure.identity import ( | ||
| DefaultAzureCredential, | ||
| DeviceCodeCredential, | ||
| InteractiveBrowserCredential | ||
| ) | ||
| from azure.core.exceptions import ClientAuthenticationError | ||
|
|
||
| # Mapping of auth types to credential classes | ||
| credential_map = { | ||
| "default": DefaultAzureCredential, | ||
| "devicecode": DeviceCodeCredential, | ||
| "interactive": InteractiveBrowserCredential, | ||
| } | ||
|
|
||
| credential_class = credential_map[auth_type] | ||
|
|
||
| try: | ||
| credential = credential_class() | ||
| token = credential.get_token("https://database.windows.net/.default").token | ||
| return AADAuth.get_token_struct(token) | ||
| except ClientAuthenticationError as e: | ||
| # Re-raise with more specific context about Azure AD authentication failure | ||
| raise RuntimeError( | ||
| f"Azure AD authentication failed for {credential_class.__name__}: {e}. " | ||
| f"This could be due to invalid credentials, missing environment variables, " | ||
| f"user cancellation, network issues, or unsupported configuration." | ||
| ) from e | ||
| except Exception as e: | ||
| # Catch any other unexpected exceptions | ||
| raise RuntimeError(f"Failed to create {credential_class.__name__}: {e}") from e | ||
|
|
||
| def process_auth_parameters(parameters: list) -> Tuple[list, Optional[str]]: | ||
jahnvi480 marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| """ | ||
| Process connection parameters and extract authentication type. | ||
|
|
||
| Args: | ||
| parameters: List of connection string parameters | ||
|
|
||
| Returns: | ||
| Tuple[list, Optional[str]]: Modified parameters and authentication type | ||
|
|
||
| Raises: | ||
| ValueError: If an invalid authentication type is provided | ||
| """ | ||
| modified_parameters = [] | ||
| auth_type = None | ||
|
|
||
| for param in parameters: | ||
| param = param.strip() | ||
| if not param: | ||
| continue | ||
|
|
||
| if "=" not in param: | ||
| modified_parameters.append(param) | ||
| continue | ||
|
|
||
| key, value = param.split("=", 1) | ||
| key_lower = key.lower() | ||
| value_lower = value.lower() | ||
|
|
||
jahnvi480 marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| if key_lower == "authentication": | ||
| # Check for supported authentication types and set auth_type accordingly | ||
| if value_lower == AuthType.INTERACTIVE.value: | ||
| # Interactive authentication (browser-based); only append parameter for non-Windows | ||
| if platform.system().lower() == "windows": | ||
| continue # Skip adding this parameter for Windows | ||
| auth_type = "interactive" | ||
| elif value_lower == AuthType.DEVICE_CODE.value: | ||
| # Device code authentication (for devices without browser) | ||
| auth_type = "devicecode" | ||
jahnvi480 marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| elif value_lower == AuthType.DEFAULT.value: | ||
| # Default authentication (uses DefaultAzureCredential) | ||
| auth_type = "default" | ||
| modified_parameters.append(param) | ||
|
|
||
| return modified_parameters, auth_type | ||
|
|
||
| def remove_sensitive_params(parameters: list) -> list: | ||
| """Remove sensitive parameters from connection string""" | ||
| exclude_keys = [ | ||
| "uid=", "pwd=", "encrypt=", "trustservercertificate=", "authentication=" | ||
| ] | ||
| return [ | ||
| param for param in parameters | ||
| if not any(param.lower().startswith(exclude) for exclude in exclude_keys) | ||
| ] | ||
|
|
||
| def get_auth_token(auth_type: str) -> Optional[bytes]: | ||
| """Get authentication token based on auth type""" | ||
| if not auth_type: | ||
| return None | ||
|
|
||
| # Handle platform-specific logic for interactive auth | ||
| if auth_type == "interactive" and platform.system().lower() == "windows": | ||
| return None # Let Windows handle AADInteractive natively | ||
|
|
||
| try: | ||
| return AADAuth.get_token(auth_type) | ||
| except (ValueError, RuntimeError): | ||
| return None | ||
|
|
||
| def process_connection_string(connection_string: str) -> Tuple[str, Optional[Dict]]: | ||
| """ | ||
| Process connection string and handle authentication. | ||
|
|
||
| Args: | ||
| connection_string: The connection string to process | ||
|
|
||
| Returns: | ||
| Tuple[str, Optional[Dict]]: Processed connection string and attrs_before dict if needed | ||
|
|
||
| Raises: | ||
| ValueError: If the connection string is invalid or empty | ||
| """ | ||
| # Check type first | ||
| if not isinstance(connection_string, str): | ||
| raise ValueError("Connection string must be a string") | ||
|
|
||
| # Then check if empty | ||
| if not connection_string: | ||
| raise ValueError("Connection string cannot be empty") | ||
|
|
||
| parameters = connection_string.split(";") | ||
jahnvi480 marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
|
||
| # Validate that there's at least one valid parameter | ||
| if not any('=' in param for param in parameters): | ||
| raise ValueError("Invalid connection string format") | ||
|
|
||
| modified_parameters, auth_type = process_auth_parameters(parameters) | ||
|
|
||
| if auth_type: | ||
| modified_parameters = remove_sensitive_params(modified_parameters) | ||
| token_struct = get_auth_token(auth_type) | ||
| if token_struct: | ||
| return ";".join(modified_parameters) + ";", {1256: token_struct} | ||
|
|
||
| return ";".join(modified_parameters) + ";", None | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.