From a4b9b7faa26d38ede72af348cc2525f897ba104d Mon Sep 17 00:00:00 2001 From: Joostlek Date: Mon, 8 May 2023 16:24:59 +0200 Subject: [PATCH] Add authentication support --- README.md | 2 +- examples/example.py | 2 +- src/python_opensky/exceptions.py | 4 ++ src/python_opensky/opensky.py | 44 ++++++++++++-- tests/ruff.toml | 1 + tests/test_states.py | 101 +++++++++++++++++++++++++++---- 6 files changed, 136 insertions(+), 18 deletions(-) diff --git a/README.md b/README.md index 0705f4f..5e3c68b 100644 --- a/README.md +++ b/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) diff --git a/examples/example.py b/examples/example.py index 63325b1..0b4bc6f 100644 --- a/examples/example.py +++ b/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) diff --git a/src/python_opensky/exceptions.py b/src/python_opensky/exceptions.py index 087964c..d68e923 100644 --- a/src/python_opensky/exceptions.py +++ b/src/python_opensky/exceptions.py @@ -11,3 +11,7 @@ class OpenSkyConnectionError(OpenSkyError): class OpenSkyCoordinateError(OpenSkyError): """OpenSky coordinate exception.""" + + +class OpenSkyUnauthenticatedError(OpenSkyError): + """OpenSky unauthenticated exception.""" diff --git a/src/python_opensky/opensky.py b/src/python_opensky/opensky.py index 8b3be29..2215aa5 100644 --- a/src/python_opensky/opensky.py +++ b/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.""" diff --git a/tests/ruff.toml b/tests/ruff.toml index b346c56..ea920d3 100644 --- a/tests/ruff.toml +++ b/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... ] diff --git a/tests/test_states.py b/tests/test_states.py index e9e24c8..7a6a811 100644 --- a/tests/test_states.py +++ b/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,7 +151,7 @@ async def test_new_session( ) async with OpenSky() as opensky: assert not opensky.session - await opensky.states() + await opensky.get_states() assert opensky.session @@ -134,7 +159,7 @@ 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()