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