diff --git a/src/datacustomcode/auth.py b/src/datacustomcode/auth.py index 6636379..ca38001 100644 --- a/src/datacustomcode/auth.py +++ b/src/datacustomcode/auth.py @@ -12,13 +12,21 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import base64 +import hashlib +import hmac import http.server import queue +import secrets import socketserver import threading import time from typing import Any -from urllib.parse import parse_qs, urlparse +from urllib.parse import ( + parse_qs, + urlencode, + urlparse, +) import webbrowser import click @@ -28,16 +36,35 @@ class OAuthCallbackHandler(http.server.SimpleHTTPRequestHandler): """HTTP request handler to capture OAuth callback.""" - def __init__(self, *args, auth_code_queue=None, **kwargs): + def __init__(self, *args, auth_code_queue=None, expected_state=None, **kwargs): self.auth_code_queue = auth_code_queue + self.expected_state = expected_state super().__init__(*args, **kwargs) + def _state_matches(self, query_params: dict) -> bool: + """Check the callback's state parameter against the expected value.""" + received = query_params.get("state", [""])[0] + if not self.expected_state or not received: + return False + return hmac.compare_digest(received, self.expected_state) + def do_GET(self): - """Handle GET request from OAuth callback.""" + """Handle GET request from OAuth callback. + + Callbacks whose ``state`` does not match are rejected without touching + the queue, so a forged request can neither inject an authorization code + nor abort the legitimate login. + """ parsed_path = urlparse(self.path) query_params = parse_qs(parsed_path.query) - if "code" in query_params: + is_callback = "code" in query_params or "error" in query_params + if is_callback and not self._state_matches(query_params): + self.send_response(400) + self.send_header("Content-type", "text/html") + self.end_headers() + self.wfile.write(b"

Invalid callback

") + elif "code" in query_params: auth_code = query_params["code"][0] self.auth_code_queue.put(auth_code) self.send_response(200) @@ -70,14 +97,27 @@ def log_message(self, format, *args): """Suppress default logging.""" +def _generate_pkce_pair() -> tuple[str, str]: + """Generate a PKCE (RFC 7636) code verifier and its S256 code challenge. + + Returns: + Tuple of (code_verifier, code_challenge) + """ + code_verifier = secrets.token_urlsafe(64) + digest = hashlib.sha256(code_verifier.encode("ascii")).digest() + code_challenge = base64.urlsafe_b64encode(digest).rstrip(b"=").decode("ascii") + return code_verifier, code_challenge + + def _run_oauth_callback_server( - redirect_uri: str, auth_code_queue: "queue.Queue[str]" + redirect_uri: str, auth_code_queue: "queue.Queue[str]", expected_state: str ) -> tuple[socketserver.TCPServer, int]: """Start a local HTTP server to catch OAuth callback. Args: redirect_uri: The redirect URI configured in the OAuth app auth_code_queue: Queue to put the authorization code in + expected_state: OAuth state value the callback must echo back Returns: Tuple of (server instance, actual port number) @@ -90,7 +130,12 @@ def _run_oauth_callback_server( # Create a custom handler factory def handler_factory(*args, **kwargs): - return OAuthCallbackHandler(*args, auth_code_queue=auth_code_queue, **kwargs) + return OAuthCallbackHandler( + *args, + auth_code_queue=auth_code_queue, + expected_state=expected_state, + **kwargs, + ) server = socketserver.TCPServer((host, port), handler_factory) server.allow_reuse_address = True @@ -113,6 +158,7 @@ def _exchange_code_for_tokens( client_secret: str, redirect_uri: str, auth_code: str, + code_verifier: str, ) -> Any: """Exchange authorization code for access and refresh tokens. @@ -122,6 +168,7 @@ def _exchange_code_for_tokens( client_secret: OAuth client secret redirect_uri: Redirect URI used in authorization auth_code: Authorization code from callback + code_verifier: PKCE code verifier matching the authorize request Returns: Dictionary containing access_token and refresh_token @@ -136,6 +183,7 @@ def _exchange_code_for_tokens( "client_id": client_id, "client_secret": client_secret, "redirect_uri": redirect_uri, + "code_verifier": code_verifier, } try: @@ -168,17 +216,28 @@ def do_oauth_browser_flow( # Create queue for communication between server and main thread auth_code_queue: queue.Queue[str] = queue.Queue() + # CSRF protection (RFC 6749 10.12) and PKCE (RFC 7636 / RFC 8252) + state = secrets.token_urlsafe(32) + code_verifier, code_challenge = _generate_pkce_pair() + # Start callback server click.echo(f"\nStarting local callback server on {redirect_uri}...") - server, _actual_port = _run_oauth_callback_server(redirect_uri, auth_code_queue) + server, _actual_port = _run_oauth_callback_server( + redirect_uri, auth_code_queue, state + ) # Build authorization URL with final redirect_uri - auth_url = ( - f"{login_url.rstrip('/')}/services/oauth2/authorize" - f"?response_type=code" - f"&client_id={client_id}" - f"&redirect_uri={redirect_uri}" + query = urlencode( + { + "response_type": "code", + "client_id": client_id, + "redirect_uri": redirect_uri, + "state": state, + "code_challenge": code_challenge, + "code_challenge_method": "S256", + } ) + auth_url = f"{login_url.rstrip('/')}/services/oauth2/authorize?{query}" # Open browser click.echo("Opening browser for authentication...") @@ -208,7 +267,7 @@ def do_oauth_browser_flow( # Exchange code for tokens click.echo("Exchanging authorization code for tokens...") token_response = _exchange_code_for_tokens( - login_url, client_id, client_secret, redirect_uri, auth_code + login_url, client_id, client_secret, redirect_uri, auth_code, code_verifier ) refresh_token = token_response.get("refresh_token") diff --git a/tests/test_auth.py b/tests/test_auth.py index 974b3c0..8095963 100644 --- a/tests/test_auth.py +++ b/tests/test_auth.py @@ -14,12 +14,16 @@ # limitations under the License. from __future__ import annotations +import base64 +import hashlib import queue from unittest.mock import ( + ANY, MagicMock, Mock, patch, ) +from urllib.parse import parse_qs, urlparse import click import pytest @@ -28,6 +32,7 @@ from datacustomcode.auth import ( OAuthCallbackHandler, _exchange_code_for_tokens, + _generate_pkce_pair, _run_oauth_callback_server, configure_oauth_tokens, do_oauth_browser_flow, @@ -48,10 +53,11 @@ def test_do_get_with_code(self, mock_init): client_address=("127.0.0.1", 5555), server=Mock(), auth_code_queue=auth_code_queue, + expected_state="expected_state_123", ) # Mock the request path with code - handler.path = "/callback?code=test_auth_code_123" + handler.path = "/callback?code=test_auth_code_123&state=expected_state_123" handler.send_response = Mock() handler.send_header = Mock() handler.end_headers = Mock() @@ -77,9 +83,13 @@ def test_do_get_with_error(self, mock_init): client_address=("127.0.0.1", 5555), server=Mock(), auth_code_queue=auth_code_queue, + expected_state="expected_state_123", ) - handler.path = "/callback?error=access_denied&error_description=User%20denied" + handler.path = ( + "/callback?error=access_denied&error_description=User%20denied" + "&state=expected_state_123" + ) handler.send_response = Mock() handler.send_header = Mock() handler.end_headers = Mock() @@ -104,6 +114,7 @@ def test_do_get_invalid_callback(self, mock_init): client_address=("127.0.0.1", 5555), server=Mock(), auth_code_queue=auth_code_queue, + expected_state="expected_state_123", ) handler.path = "/callback" @@ -118,6 +129,63 @@ def test_do_get_invalid_callback(self, mock_init): assert auth_code_queue.empty() handler.send_response.assert_called_once_with(400) + @pytest.mark.parametrize( + "path", + [ + "/callback?code=attacker_code", + "/callback?code=attacker_code&state=", + "/callback?code=attacker_code&state=wrong_state", + "/callback?error=access_denied&state=wrong_state", + "/callback?error=access_denied", + ], + ) + @patch("datacustomcode.auth.http.server.SimpleHTTPRequestHandler.__init__") + def test_do_get_rejects_missing_or_mismatched_state(self, mock_init, path): + """Callbacks without the expected state must not reach the queue.""" + auth_code_queue = queue.Queue() + mock_init.return_value = None + + handler = OAuthCallbackHandler( + request=Mock(), + client_address=("127.0.0.1", 5555), + server=Mock(), + auth_code_queue=auth_code_queue, + expected_state="expected_state_123", + ) + handler.path = path + handler.send_response = Mock() + handler.send_header = Mock() + handler.end_headers = Mock() + handler.wfile = Mock() + + handler.do_GET() + + assert auth_code_queue.empty() + handler.send_response.assert_called_once_with(400) + + @patch("datacustomcode.auth.http.server.SimpleHTTPRequestHandler.__init__") + def test_do_get_rejects_when_no_expected_state(self, mock_init): + """A handler without an expected state fails closed.""" + auth_code_queue = queue.Queue() + mock_init.return_value = None + + handler = OAuthCallbackHandler( + request=Mock(), + client_address=("127.0.0.1", 5555), + server=Mock(), + auth_code_queue=auth_code_queue, + ) + handler.path = "/callback?code=some_code&state=anything" + handler.send_response = Mock() + handler.send_header = Mock() + handler.end_headers = Mock() + handler.wfile = Mock() + + handler.do_GET() + + assert auth_code_queue.empty() + handler.send_response.assert_called_once_with(400) + @patch("datacustomcode.auth.http.server.SimpleHTTPRequestHandler.__init__") def test_log_message_suppressed(self, mock_init): """Test that log messages are suppressed.""" @@ -148,7 +216,9 @@ def test_run_oauth_callback_server_valid_uri( mock_server.server_address = ("localhost", 5555) mock_tcpserver.return_value = mock_server - _server, port = _run_oauth_callback_server(redirect_uri, auth_code_queue) + _server, port = _run_oauth_callback_server( + redirect_uri, auth_code_queue, "state" + ) assert port == 5555 mock_tcpserver.assert_called_once() @@ -162,7 +232,7 @@ def test_run_oauth_callback_server_invalid_uri_no_host(self): auth_code_queue = queue.Queue() with pytest.raises(ValueError, match="Invalid redirect URI"): - _run_oauth_callback_server(redirect_uri, auth_code_queue) + _run_oauth_callback_server(redirect_uri, auth_code_queue, "state") def test_run_oauth_callback_server_invalid_uri_no_port(self): """Test starting callback server with invalid URI (no port).""" @@ -170,7 +240,7 @@ def test_run_oauth_callback_server_invalid_uri_no_port(self): auth_code_queue = queue.Queue() with pytest.raises(ValueError, match="Invalid redirect URI"): - _run_oauth_callback_server(redirect_uri, auth_code_queue) + _run_oauth_callback_server(redirect_uri, auth_code_queue, "state") @patch("datacustomcode.auth.socketserver.TCPServer") @patch("datacustomcode.auth.threading.Thread") @@ -186,7 +256,9 @@ def test_run_oauth_callback_server_different_port( mock_server.server_address = ("localhost", 8080) mock_tcpserver.return_value = mock_server - _server, port = _run_oauth_callback_server(redirect_uri, auth_code_queue) + _server, port = _run_oauth_callback_server( + redirect_uri, auth_code_queue, "state" + ) assert port == 8080 mock_tcpserver.assert_called_once() @@ -195,6 +267,28 @@ def test_run_oauth_callback_server_different_port( assert mock_server.allow_reuse_address is True +class TestGeneratePkcePair: + """Test suite for _generate_pkce_pair.""" + + def test_verifier_and_challenge(self): + code_verifier, code_challenge = _generate_pkce_pair() + + # RFC 7636 section 4.1: 43-128 unreserved characters + assert 43 <= len(code_verifier) <= 128 + expected = ( + base64.urlsafe_b64encode( + hashlib.sha256(code_verifier.encode("ascii")).digest() + ) + .rstrip(b"=") + .decode("ascii") + ) + assert code_challenge == expected + assert "=" not in code_challenge + + def test_pairs_are_unique(self): + assert _generate_pkce_pair()[0] != _generate_pkce_pair()[0] + + class TestExchangeCodeForTokens: """Test suite for _exchange_code_for_tokens.""" @@ -216,6 +310,7 @@ def test_exchange_code_for_tokens_success(self, mock_post): client_secret="test_client_secret", redirect_uri="http://localhost:5555/callback", auth_code="test_auth_code", + code_verifier="test_code_verifier", ) assert result["access_token"] == "test_access_token" @@ -228,6 +323,7 @@ def test_exchange_code_for_tokens_success(self, mock_post): assert call_args[1]["data"]["client_id"] == "test_client_id" assert call_args[1]["data"]["client_secret"] == "test_client_secret" assert call_args[1]["data"]["redirect_uri"] == "http://localhost:5555/callback" + assert call_args[1]["data"]["code_verifier"] == "test_code_verifier" @patch("datacustomcode.auth.requests.post") def test_exchange_code_for_tokens_http_error(self, mock_post): @@ -241,6 +337,7 @@ def test_exchange_code_for_tokens_http_error(self, mock_post): client_secret="test_client_secret", redirect_uri="http://localhost:5555/callback", auth_code="test_auth_code", + code_verifier="test_code_verifier", ) @patch("datacustomcode.auth.requests.post") @@ -255,6 +352,7 @@ def test_exchange_code_for_tokens_connection_error(self, mock_post): client_secret="test_client_secret", redirect_uri="http://localhost:5555/callback", auth_code="test_auth_code", + code_verifier="test_code_verifier", ) @patch("datacustomcode.auth.requests.post") @@ -269,6 +367,7 @@ def test_exchange_code_for_tokens_timeout(self, mock_post): client_secret="test_client_secret", redirect_uri="http://localhost:5555/callback", auth_code="test_auth_code", + code_verifier="test_code_verifier", ) @@ -288,7 +387,7 @@ def test_do_oauth_browser_flow_success( auth_code_queue = queue.Queue() auth_code_queue.put("test_auth_code") - def server_factory(redirect_uri, queue): + def server_factory(redirect_uri, queue, expected_state): return (mock_server_instance, 5555) mock_server.side_effect = server_factory @@ -315,9 +414,34 @@ def server_factory(redirect_uri, queue): "test_client_secret", "http://localhost:5555/callback", "test_auth_code", + ANY, ) mock_server_instance.shutdown.assert_called() + # Authorize URL carries state + S256 PKCE challenge + auth_url = mock_browser.call_args[0][0] + params = parse_qs(urlparse(auth_url).query) + assert params["response_type"] == ["code"] + assert params["client_id"] == ["test_client_id"] + assert params["redirect_uri"] == ["http://localhost:5555/callback"] + assert params["code_challenge_method"] == ["S256"] + + # The callback server expects the same state sent to the browser + server_state = mock_server.call_args[0][2] + assert params["state"] == [server_state] + assert len(server_state) >= 32 + + # The verifier sent at token exchange matches the challenge + code_verifier = mock_exchange.call_args[0][5] + expected_challenge = ( + base64.urlsafe_b64encode( + hashlib.sha256(code_verifier.encode("ascii")).digest() + ) + .rstrip(b"=") + .decode("ascii") + ) + assert params["code_challenge"] == [expected_challenge] + @patch("datacustomcode.auth.webbrowser.open") @patch("datacustomcode.auth._run_oauth_callback_server") @patch("datacustomcode.auth.click.echo") @@ -325,7 +449,7 @@ def test_do_oauth_browser_flow_timeout(self, mock_echo, mock_server, mock_browse """Test OAuth browser flow with timeout.""" mock_server_instance = MagicMock() - def server_factory(redirect_uri, queue): + def server_factory(redirect_uri, queue, expected_state): return (mock_server_instance, 5555) mock_server.side_effect = server_factory @@ -356,7 +480,7 @@ def test_do_oauth_browser_flow_oauth_error( """Test OAuth browser flow with OAuth error response.""" mock_server_instance = MagicMock() - def server_factory(redirect_uri, queue): + def server_factory(redirect_uri, queue, expected_state): return (mock_server_instance, 5555) mock_server.side_effect = server_factory @@ -388,7 +512,7 @@ def test_do_oauth_browser_flow_no_refresh_token( """Test OAuth browser flow when no refresh token is returned.""" mock_server_instance = MagicMock() - def server_factory(redirect_uri, queue): + def server_factory(redirect_uri, queue, expected_state): return (mock_server_instance, 5555) mock_server.side_effect = server_factory