[client] add JWT management methods

Author: capcom6

android_sms_gateway/client.py

@@ -16,15 +16,17 @@ class BaseClient(abc.ABC):
     def __init__(
         self,
         login: t.Optional[str],
-        password_or_token: str,
+        password: str,
         *,
         base_url: str = DEFAULT_URL,
         encryptor: t.Optional[BaseEncryptor] = None,
     ) -> None:
-        if login and password_or_token:
-            auth_header = f"Basic {base64.b64encode(f'{login}:{password_or_token}'.encode()).decode()}"
-        elif password_or_token:
-            auth_header = f"Bearer {password_or_token}"
+        if login and password:
+            auth_header = (
+                f"Basic {base64.b64encode(f'{login}:{password}'.encode()).decode()}"
+            )
+        elif password:
+            auth_header = f"Bearer {password}"
         else:
             raise ValueError("Either login and password or token must be provided")
 
@@ -92,15 +94,15 @@ class APIClient(BaseClient):
     def __init__(
         self,
         login: t.Optional[str],
-        password_or_token: str,
+        password: str,
         *,
         base_url: str = DEFAULT_URL,
         encryptor: t.Optional[BaseEncryptor] = None,
         http: t.Optional[http.HttpClient] = None,
     ) -> None:
         super().__init__(
             login=login,
-            password_or_token=password_or_token,
+            password=password,
             base_url=base_url,
             encryptor=encryptor,
         )
@@ -226,20 +228,55 @@ def health_check(self) -> dict:
 
         return self.http.get(f"{self.base_url}/health", headers=self.headers)
 
+    def generate_token(
+        self, token_request: domain.TokenRequest
+    ) -> domain.TokenResponse:
+        """
+        Generates a new JWT token with specified scopes and TTL.
+
+        Args:
+            token_request: The token request containing scopes and optional TTL.
+
+        Returns:
+            The generated token response.
+        """
+        if self.http is None:
+            raise ValueError("HTTP client not initialized")
+
+        return domain.TokenResponse.from_dict(
+            self.http.post(
+                f"{self.base_url}/auth/token",
+                payload=token_request.asdict(),
+                headers=self.headers,
+            )
+        )
+
+    def revoke_token(self, jti: str) -> None:
+        """
+        Revokes a JWT token with the specified JTI (token ID).
+
+        Args:
+            jti: The JTI (token ID) of the token to revoke.
+        """
+        if self.http is None:
+            raise ValueError("HTTP client not initialized")
+
+        self.http.delete(f"{self.base_url}/auth/token/{jti}", headers=self.headers)
+
 
 class AsyncAPIClient(BaseClient):
     def __init__(
         self,
         login: t.Optional[str],
-        password_or_token: str,
+        password: str,
         *,
         base_url: str = DEFAULT_URL,
         encryptor: t.Optional[BaseEncryptor] = None,
         http_client: t.Optional[ahttp.AsyncHttpClient] = None,
     ) -> None:
         super().__init__(
             login=login,
-            password_or_token=password_or_token,
+            password=password,
             base_url=base_url,
             encryptor=encryptor,
         )
@@ -258,7 +295,7 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
         if self.default_http is None:
             return
 
-        self.default_http.__aexit__(exc_type, exc_val, exc_tb)
+        await self.default_http.__aexit__(exc_type, exc_val, exc_tb)
         self.http = self.default_http = None
 
     async def send(self, message: domain.Message) -> domain.MessageState:
@@ -366,3 +403,40 @@ async def health_check(self) -> dict:
             raise ValueError("HTTP client not initialized")
 
         return await self.http.get(f"{self.base_url}/health", headers=self.headers)
+
+    async def generate_token(
+        self, token_request: domain.TokenRequest
+    ) -> domain.TokenResponse:
+        """
+        Generates a new JWT token with specified scopes and TTL.
+
+        Args:
+            token_request: The token request containing scopes and optional TTL.
+
+        Returns:
+            The generated token response.
+        """
+        if self.http is None:
+            raise ValueError("HTTP client not initialized")
+
+        return domain.TokenResponse.from_dict(
+            await self.http.post(
+                f"{self.base_url}/auth/token",
+                payload=token_request.asdict(),
+                headers=self.headers,
+            )
+        )
+
+    async def revoke_token(self, jti: str) -> None:
+        """
+        Revokes a JWT token with the specified JTI (token ID).
+
+        Args:
+            jti: The JTI (token ID) of the token to revoke.
+        """
+        if self.http is None:
+            raise ValueError("HTTP client not initialized")
+
+        await self.http.delete(
+            f"{self.base_url}/auth/token/{jti}", headers=self.headers
+        )

android_sms_gateway/domain.py

@@ -266,3 +266,57 @@ def from_dict(cls, payload: t.Dict[str, t.Any]) -> "ErrorResponse":
             code=payload["code"],
             message=payload["message"],
         )
+
+
+@dataclasses.dataclass(frozen=True)
+class TokenRequest:
+    """Represents a request to generate a new JWT token."""
+
+    scopes: t.List[str]
+    """List of scopes for the token."""
+    ttl: t.Optional[int] = None
+    """Time to live for the token in seconds."""
+
+    def asdict(self) -> t.Dict[str, t.Any]:
+        """Returns a dictionary representation of the token request.
+
+        Returns:
+            A dictionary containing the token request data.
+        """
+        result: t.Dict[str, t.Any] = {
+            "scopes": self.scopes,
+        }
+        if self.ttl is not None:
+            result["ttl"] = self.ttl
+        return result
+
+
+@dataclasses.dataclass(frozen=True)
+class TokenResponse:
+    """Represents a response when generating a new JWT token."""
+
+    access_token: str
+    """The JWT access token."""
+    token_type: str
+    """The type of the token (e.g., 'Bearer')."""
+    id: str
+    """The unique identifier of the token (jti)."""
+    expires_at: str
+    """The expiration time of the token in ISO format."""
+
+    @classmethod
+    def from_dict(cls, payload: t.Dict[str, t.Any]) -> "TokenResponse":
+        """Creates a TokenResponse instance from a dictionary.
+
+        Args:
+            payload: A dictionary containing the token response data.
+
+        Returns:
+            A TokenResponse instance.
+        """
+        return cls(
+            access_token=payload["accessToken"],
+            token_type=payload["tokenType"],
+            id=payload["id"],
+            expires_at=payload["expiresAt"],
+        )

tests/test_jwt_auth.py

@@ -1,6 +1,9 @@
 import pytest
+from unittest.mock import AsyncMock, MagicMock
+
 from android_sms_gateway.client import APIClient, AsyncAPIClient
 from android_sms_gateway.constants import DEFAULT_URL
+from android_sms_gateway.domain import TokenRequest, TokenResponse
 from android_sms_gateway.http import RequestsHttpClient
 
 
@@ -28,7 +31,7 @@ def test_jwt_auth_initialization():
         RequestsHttpClient() as h,
         APIClient(
             login=None,
-            password_or_token="test_jwt_token",
+            password="test_jwt_token",
             base_url=DEFAULT_URL,
             http=h,
         ) as client,
@@ -58,7 +61,7 @@ def test_async_jwt_auth_initialization():
     """Test that the async client can be initialized with JWT token."""
     client = AsyncAPIClient(
         login=None,
-        password_or_token="test_jwt_token",
+        password="test_jwt_token",
         base_url=DEFAULT_URL,
     )
     # Check that the Authorization header is set correctly
@@ -83,7 +86,7 @@ def test_missing_password_error():
         ValueError,
         match="Either login and password or token must be provided",
     ):
-        APIClient(login="test_login", password_or_token="")
+        APIClient(login="test_login", password="")
 
 
 def test_async_missing_credentials_error():
@@ -93,3 +96,115 @@ def test_async_missing_credentials_error():
         match="Either login and password or token must be provided",
     ):
         AsyncAPIClient(None, "")
+
+
+def test_generate_token():
+    """Test that the client can generate a new JWT token."""
+    mock_http = MagicMock()
+    mock_http.post.return_value = {
+        "accessToken": "test_access_token",
+        "tokenType": "Bearer",
+        "id": "test_token_id",
+        "expiresAt": "2023-12-31T23:59:59Z",
+    }
+
+    token_request = TokenRequest(scopes=["sms:send", "sms:read"], ttl=3600)
+
+    with APIClient(
+        login=None,
+        password="initial_token",
+        base_url=DEFAULT_URL,
+        http=mock_http,
+    ) as client:
+        response = client.generate_token(token_request)
+
+        # Verify the response
+        assert isinstance(response, TokenResponse)
+        assert response.access_token == "test_access_token"
+        assert response.token_type == "Bearer"
+        assert response.id == "test_token_id"
+        assert response.expires_at == "2023-12-31T23:59:59Z"
+
+        # Verify the HTTP call
+        mock_http.post.assert_called_once_with(
+            f"{DEFAULT_URL}/auth/token",
+            payload=token_request.asdict(),
+            headers=client.headers,
+        )
+
+
+def test_revoke_token():
+    """Test that the client can revoke a JWT token."""
+    mock_http = MagicMock()
+    mock_http.delete.return_value = None
+
+    with APIClient(
+        login=None,
+        password="initial_token",
+        base_url=DEFAULT_URL,
+        http=mock_http,
+    ) as client:
+        client.revoke_token("test_token_id")
+
+        # Verify the HTTP call
+        mock_http.delete.assert_called_once_with(
+            f"{DEFAULT_URL}/auth/token/test_token_id",
+            headers=client.headers,
+        )
+
+
+@pytest.mark.asyncio
+async def test_async_generate_token():
+    """Test that the async client can generate a new JWT token."""
+    mock_http = AsyncMock()
+    mock_http.post.return_value = {
+        "accessToken": "test_access_token",
+        "tokenType": "Bearer",
+        "id": "test_token_id",
+        "expiresAt": "2023-12-31T23:59:59Z",
+    }
+
+    token_request = TokenRequest(scopes=["sms:send", "sms:read"], ttl=3600)
+
+    async with AsyncAPIClient(
+        login=None,
+        password="initial_token",
+        base_url=DEFAULT_URL,
+        http_client=mock_http,
+    ) as client:
+        response = await client.generate_token(token_request)
+
+        # Verify the response
+        assert isinstance(response, TokenResponse)
+        assert response.access_token == "test_access_token"
+        assert response.token_type == "Bearer"
+        assert response.id == "test_token_id"
+        assert response.expires_at == "2023-12-31T23:59:59Z"
+
+        # Verify the HTTP call
+        mock_http.post.assert_called_once_with(
+            f"{DEFAULT_URL}/auth/token",
+            payload=token_request.asdict(),
+            headers=client.headers,
+        )
+
+
+@pytest.mark.asyncio
+async def test_async_revoke_token():
+    """Test that the async client can revoke a JWT token."""
+    mock_http = AsyncMock()
+    mock_http.delete.return_value = None
+
+    async with AsyncAPIClient(
+        login=None,
+        password="initial_token",
+        base_url=DEFAULT_URL,
+        http_client=mock_http,
+    ) as client:
+        await client.revoke_token("test_token_id")
+
+        # Verify the HTTP call
+        mock_http.delete.assert_called_once_with(
+            f"{DEFAULT_URL}/auth/token/test_token_id",
+            headers=client.headers,
+        )