Merge pull request #17 from joostlek/auth

Add authentication support

Author: joostlek

README.md

@@ -33,7 +33,7 @@ from python_opensky import OpenSky, StatesResponse
 async def main() -> None:
     """Show example of fetching all flight states."""
     async with OpenSky() as opensky:
-        states: StatesResponse = await opensky.states()
+        states: StatesResponse = await opensky.get_states()
         print(states)
 
 

examples/example.py

@@ -8,7 +8,7 @@
 async def main() -> None:
     """Show example of fetching flight states from OpenSky."""
     async with OpenSky() as opensky:
-        states: StatesResponse = await opensky.states()
+        states: StatesResponse = await opensky.get_states()
         print(states)
 
 

src/python_opensky/exceptions.py

@@ -11,3 +11,7 @@ class OpenSkyConnectionError(OpenSkyError):
 
 class OpenSkyCoordinateError(OpenSkyError):
     """OpenSky coordinate exception."""
+
+
+class OpenSkyUnauthenticatedError(OpenSkyError):
+    """OpenSky unauthenticated exception."""

src/python_opensky/opensky.py

@@ -9,12 +9,16 @@
 from typing import Any, cast
 
 import async_timeout
-from aiohttp import ClientError, ClientResponseError, ClientSession
+from aiohttp import BasicAuth, ClientError, ClientResponseError, ClientSession
 from aiohttp.hdrs import METH_GET
 from yarl import URL
 
 from .const import MAX_LATITUDE, MAX_LONGITUDE, MIN_LATITUDE, MIN_LONGITUDE
-from .exceptions import OpenSkyConnectionError, OpenSkyError
+from .exceptions import (
+    OpenSkyConnectionError,
+    OpenSkyError,
+    OpenSkyUnauthenticatedError,
+)
 from .models import BoundingBox, StatesResponse
 
 
@@ -24,11 +28,22 @@ class OpenSky:
 
     session: ClientSession | None = None
     request_timeout: int = 10
-    api_host: str = "python_opensky-network.org"
+    api_host: str = "opensky-network.org"
     opensky_credits: int = 400
     timezone = timezone.utc
     _close_session: bool = False
     _credit_usage: dict[datetime, int] = field(default_factory=dict)
+    _auth: BasicAuth | None = None
+    _contributing_user: bool = False
+
+    def authenticate(self, auth: BasicAuth, *, contributing_user: bool = False) -> None:
+        """Authenticate the user."""
+        self._auth = auth
+        self._contributing_user = contributing_user
+        if contributing_user:
+            self.opensky_credits = 8000
+        else:
+            self.opensky_credits = 4000
 
     async def _request(
         self,
@@ -79,6 +94,7 @@ async def _request(
                 response = await self.session.request(
                     METH_GET,
                     url.with_query(data),
+                    auth=self._auth,
                     headers=headers,
                 )
                 response.raise_for_status()
@@ -105,7 +121,10 @@ async def _request(
 
         return cast(dict[str, Any], await response.json())
 
-    async def states(self, bounding_box: BoundingBox | None = None) -> StatesResponse:
+    async def get_states(
+        self,
+        bounding_box: BoundingBox | None = None,
+    ) -> StatesResponse:
         """Retrieve state vectors for a given time."""
         credit_cost = 4
         params = {
@@ -132,6 +151,23 @@ async def states(self, bounding_box: BoundingBox | None = None) -> StatesRespons
 
         return StatesResponse.parse_obj(data)
 
+    async def get_own_states(self, time: int = 0) -> StatesResponse:
+        """Retrieve state vectors from your own sensors."""
+        if not self._auth:
+            raise OpenSkyUnauthenticatedError
+        params = {
+            "time": time,
+        }
+
+        data = await self._request("states/own", data=params)
+
+        data = {
+            **data,
+            "states": [self._convert_state(state) for state in data["states"]],
+        }
+
+        return StatesResponse.parse_obj(data)
+
     @staticmethod
     def calculate_credit_costs(bounding_box: BoundingBox) -> int:
         """Calculate the amount of credits a request costs."""

tests/ruff.toml

@@ -7,6 +7,7 @@ extend-select = [
 
 extend-ignore = [
   "S101", # Use of assert detected. As these are tests...
+  "S106", # Detection of passwords...
   "SLF001", # Tests will access private/protected members...
   "TCH002", # pytest doesn't like this one...
 ]

tests/test_states.py

@@ -3,7 +3,8 @@
 
 import aiohttp
 import pytest
-from aiohttp import ClientError
+from aiohttp import BasicAuth, ClientError
+from aiohttp.web_request import BaseRequest
 from aresponses import Response, ResponsesMockServer
 
 from python_opensky import (
@@ -15,10 +16,11 @@
     PositionSource,
     StatesResponse,
 )
+from python_opensky.exceptions import OpenSkyUnauthenticatedError
 
 from . import load_fixture
 
-OPENSKY_URL = "python_opensky-network.org"
+OPENSKY_URL = "opensky-network.org"
 
 
 async def test_states(
@@ -37,7 +39,7 @@ async def test_states(
     )
     async with aiohttp.ClientSession() as session:
         opensky = OpenSky(session=session)
-        response: StatesResponse = await opensky.states()
+        response: StatesResponse = await opensky.get_states()
         assert len(response.states) == 4
         assert response.time == 1683488744
         first_aircraft = response.states[0]
@@ -62,6 +64,29 @@ async def test_states(
         await opensky.close()
 
 
+async def test_own_states(
+    aresponses: ResponsesMockServer,
+) -> None:
+    """Test retrieving own states."""
+    aresponses.add(
+        OPENSKY_URL,
+        "/api/states/own",
+        "GET",
+        aresponses.Response(
+            status=200,
+            headers={"Content-Type": "application/json"},
+            text=load_fixture("states.json"),
+        ),
+    )
+    async with aiohttp.ClientSession() as session:
+        opensky = OpenSky(session=session)
+        opensky.authenticate(BasicAuth(login="test", password="test"))
+        response: StatesResponse = await opensky.get_own_states()
+        assert len(response.states) == 4
+        assert opensky.remaining_credits() == opensky.opensky_credits
+        await opensky.close()
+
+
 async def test_states_with_bounding_box(
     aresponses: ResponsesMockServer,
 ) -> None:
@@ -85,7 +110,7 @@ async def test_states_with_bounding_box(
             min_longitude=0,
             max_longitude=0,
         )
-        await opensky.states(bounding_box=bounding_box)
+        await opensky.get_states(bounding_box=bounding_box)
         await opensky.close()
 
 
@@ -105,8 +130,8 @@ async def test_credit_usage(
     )
     async with aiohttp.ClientSession() as session:
         opensky = OpenSky(session=session)
-        await opensky.states()
-        assert opensky.remaining_credits() == 396
+        await opensky.get_states()
+        assert opensky.remaining_credits() == opensky.opensky_credits - 4
         await opensky.close()
 
 
@@ -126,15 +151,15 @@ async def test_new_session(
     )
     async with OpenSky() as opensky:
         assert not opensky.session
-        await opensky.states()
+        await opensky.get_states()
         assert opensky.session
 
 
 async def test_timeout(aresponses: ResponsesMockServer) -> None:
     """Test request timeout."""
 
     # Faking a timeout by sleeping
-    async def response_handler(_: aiohttp.ClientResponse) -> Response:
+    async def response_handler(_: BaseRequest) -> Response:
         """Response handler for this test."""
         await asyncio.sleep(2)
         return aresponses.Response(body="Goodmorning!")
@@ -149,14 +174,57 @@ async def response_handler(_: aiohttp.ClientResponse) -> Response:
     async with aiohttp.ClientSession() as session:
         opensky = OpenSky(session=session, request_timeout=1)
         with pytest.raises(OpenSkyConnectionError):
-            assert await opensky.states()
+            assert await opensky.get_states()
+        await opensky.close()
+
+
+async def test_auth(aresponses: ResponsesMockServer) -> None:
+    """Test request authentication."""
+
+    def response_handler(request: BaseRequest) -> Response:
+        """Response handler for this test."""
+        assert request.headers
+        assert request.headers["Authorization"]
+        assert request.headers["Authorization"] == "Basic dGVzdDp0ZXN0"
+        return aresponses.Response(
+            status=200,
+            headers={"Content-Type": "application/json"},
+            text=load_fixture("states.json"),
+        )
+
+    aresponses.add(
+        OPENSKY_URL,
+        "/api/states/all",
+        "GET",
+        response_handler,
+    )
+
+    async with aiohttp.ClientSession() as session:
+        opensky = OpenSky(session=session)
+        opensky.authenticate(BasicAuth(login="test", password="test"))
+        await opensky.get_states()
+        await opensky.close()
+
+
+async def test_user_credits() -> None:
+    """Test authenticated user credits."""
+    async with aiohttp.ClientSession() as session:
+        opensky = OpenSky(session=session)
+        assert opensky.opensky_credits == 400
+        opensky.authenticate(BasicAuth(login="test", password="test"))
+        assert opensky.opensky_credits == 4000
+        opensky.authenticate(
+            BasicAuth(login="test", password="test"),
+            contributing_user=True,
+        )
+        assert opensky.opensky_credits == 8000
         await opensky.close()
 
 
 async def test_request_error(aresponses: ResponsesMockServer) -> None:
     """Test request error."""
 
-    async def response_handler(_: aiohttp.ClientResponse) -> Response:
+    async def response_handler(_: BaseRequest) -> Response:
         """Response handler for this test."""
         raise ClientError
 
@@ -170,7 +238,7 @@ async def response_handler(_: aiohttp.ClientResponse) -> Response:
     async with aiohttp.ClientSession() as session:
         opensky = OpenSky(session=session)
         with pytest.raises(OpenSkyConnectionError):
-            assert await opensky.states()
+            assert await opensky.get_states()
         await opensky.close()
 
 
@@ -192,7 +260,16 @@ async def test_unexpected_server_response(
     async with aiohttp.ClientSession() as session:
         opensky = OpenSky(session=session)
         with pytest.raises(OpenSkyError):
-            assert await opensky.states()
+            assert await opensky.get_states()
+        await opensky.close()
+
+
+async def test_unauthenticated_own_states() -> None:
+    """Test unauthenticated access to own states."""
+    async with aiohttp.ClientSession() as session:
+        opensky = OpenSky(session=session)
+        with pytest.raises(OpenSkyUnauthenticatedError):
+            assert await opensky.get_own_states()
         await opensky.close()