Skip to content

Commit

Permalink
Asyncio support
Browse files Browse the repository at this point in the history
  • Loading branch information
Kylmakalle committed Nov 17, 2023
1 parent 41b520d commit 45e3cc6
Show file tree
Hide file tree
Showing 7 changed files with 138 additions and 7 deletions.
6 changes: 6 additions & 0 deletions Readme.md
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,12 @@ device_check = DeviceCheck(
)
```

### Asyncio setup
```python
from devicecheck.asyncio import AsyncioDeviceCheck
```
The rest will be the same, except for network methods must be `await`'ed

### Validate device

```python
Expand Down
10 changes: 7 additions & 3 deletions devicecheck/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,10 @@
https://github.com/Kylmakalle/devicecheck
"""
__version__ = "1.2.2"
__version__ = "1.3.0"
__author__ = 'Sergey Akentev (@Kylmakalle)'
__license__ = 'MIT'
__copyright__ = 'Copyright 2021 Sergey Akentev'
__copyright__ = 'Copyright 2023 Sergey Akentev'

import logging
import os
Expand Down Expand Up @@ -207,7 +207,11 @@ def __init__(self, team_id: str, bundle_id: str, key_id: str, private_key: [str,
log.warning("Using Development environment. Remember to set dev_environment=False in production!")
self.retry_wrong_env_request = retry_wrong_env_request
self.raise_on_error = raise_on_error
self._session = requests.Session()
self._session = self._make_session()

@staticmethod
def _make_session() -> requests.Session:
return requests.Session()

def generate_token(self, valid_time: int = 500, force_refresh=False):
"""
Expand Down
93 changes: 93 additions & 0 deletions devicecheck/asyncio.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,93 @@
import aiohttp
import certifi
import ssl
from . import DeviceCheck, get_timestamp_milliseconds, get_transaction_id, parse_apple_response, HttpAppleResponse, log, DataAppleResponse


class AsyncioDeviceCheck(DeviceCheck):

@staticmethod
def _make_session() -> aiohttp.ClientSession:
return aiohttp.ClientSession(
connector=aiohttp.TCPConnector(
ssl=ssl.create_default_context(cafile=certifi.where())
),
)

async def validate_device_token(self, token: str, *args, **kwargs) -> HttpAppleResponse:
"""
Validate a device with it's token
https://developer.apple.com/documentation/devicecheck/accessing_and_modifying_per-device_data#2929855
:param token: Base 64–encoded representation of encrypted device information
:param args: Additional args for requests module
:param kwargs: Additional kwargs for requests module
:return:
"""
endpoint = 'v1/validate_device_token'
payload = {
'timestamp': get_timestamp_milliseconds(),
'transaction_id': get_transaction_id(),
'device_token': token
}

result = await self._request(self.get_request_url(), endpoint, payload, *args, **kwargs)
return parse_apple_response(result, self.raise_on_error)

async def query_two_bits(self, token: str, *args, **kwargs) -> DataAppleResponse:
"""
Query two bits of device data
https://developer.apple.com/documentation/devicecheck/accessing_and_modifying_per-device_data#2910405
:param token: Base 64–encoded representation of encrypted device information
:param args: Additional args for requests module
:param kwargs: Additional kwargs for requests module
:return:
"""
endpoint = 'v1/query_two_bits'
payload = {
'timestamp': get_timestamp_milliseconds(),
'transaction_id': get_transaction_id(),
'device_token': token
}
result = await self._request(self.get_request_url(), endpoint, payload, *args, **kwargs)
return parse_apple_response(result, self.raise_on_error)

async def update_two_bits(self, token: str, bit_0: [bool, int] = None, bit_1: [bool, int] = None, *args,
**kwargs) -> HttpAppleResponse:
"""
Update bit(s) of a device data
https://developer.apple.com/documentation/devicecheck/accessing_and_modifying_per-device_data#2910405
:param token: Base 64–encoded representation of encrypted device information
:param bit_0: (Optional) First bit of data `bool`
:param bit_1: (Optional) Second bit of data `bool`
:param args: Additional args for requests module
:param kwargs: Additional kwargs for requests module
:return:
"""
endpoint = 'v1/update_two_bits'
payload = {
'timestamp': get_timestamp_milliseconds(),
'transaction_id': get_transaction_id(),
'device_token': token,
}

if bit_0 is not None:
payload['bit0'] = bool(bit_0)

if bit_1 is not None:
payload['bit1'] = bool(bit_1)

result = await self._request(self.get_request_url(), endpoint, payload, *args, **kwargs)
return parse_apple_response(result, self.raise_on_error)

async def _request(self, url, endpoint, payload, retrying_env=False, *args, **kwargs):
log.debug(f'Sending request to {url}/{endpoint} with data {payload}')

async with self._session.post(f"{url}/{endpoint}", json=payload, headers={"authorization": f"Bearer {self.generate_token()}"}) as response:
result_text = await response.text()
log.debug(f"Response: {response.status} {result_text}")

if not retrying_env and self.retry_wrong_env_request and response.status != 200:
log.info(f"Retrying request on {'production' if not self.dev_environment else 'development'} server")
return await self._request(self.get_request_url(not self.dev_environment), endpoint, payload, retrying_env=True, *args, **kwargs)

return response.status, result_text
16 changes: 13 additions & 3 deletions devicecheck/decorators.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
from functools import wraps

from . import DeviceCheck, AppleException
from .asyncio import AsyncioDeviceCheck

log = logging.getLogger('devicecheck:decorator')
logging.basicConfig()
Expand Down Expand Up @@ -63,13 +64,22 @@ def _is_valid_device(device_check_instance: DeviceCheck, token: str):
return False # pragma: no cover


def async_validate_device(device_check_instance: DeviceCheck,
async def _is_valid_device_async(device_check_instance: AsyncioDeviceCheck, token: str):
if MOCK_DEVICE_CHECK_DECORATOR_TOKEN:
return token == MOCK_DEVICE_CHECK_DECORATOR_TOKEN
try: # pragma: no cover
return await device_check_instance.validate_device_token(token).is_ok # pragma: no cover
except AppleException: # pragma: no cover
return False # pragma: no cover


def async_validate_device(device_check_instance: AsyncioDeviceCheck,
framework: [DCSupportedAsyncFrameworks, str] = None,
on_invalid_token=('Invalid device token', 403)):
"""
Async Decorator that validates device token provided in `Device-Token` header
or `device_token`/`deviceToken` key in json body.
:param device_check_instance: Instance of DeviceCheck module for validating
:param device_check_instance: Instance of AsyncioDeviceCheck module for validating
:param framework: Name of used async framework for automated data extraction. Leave `None` to rely on a universal parser.
:param on_invalid_token: Object that will be returned if validation was unsuccessful
:return: on_invalid_token variable
Expand Down Expand Up @@ -104,7 +114,7 @@ async def decorated_function(*args, **kwargs):
device_token = await async_extract_device_token(request, framework)
if device_token:
try:
is_valid = _is_valid_device(device_check_instance, device_token)
is_valid = await _is_valid_device_async(device_check_instance, device_token)
except Exception as e: # pragma: no cover
log.error(f'DeviceCheck request failed. {e}') # pragma: no cover
if is_valid:
Expand Down
3 changes: 3 additions & 0 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,9 @@ def get_version(rel_path):
"pyjwt>=2.0.0",
"cryptography>=3.4.7"
],
extras_require={
'async': ['aiohttp>=3.8']
},
packages=["devicecheck"],
python_requires=">=3.6",
)
15 changes: 15 additions & 0 deletions tests/integration/asynciodevicecheck_mock.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
from devicecheck.asyncio import AsyncioDeviceCheck

PRIVATE_KEY = """
-----BEGIN PRIVATE KEY-----
TEST
-----END PRIVATE KEY-----
"""

device_check = AsyncioDeviceCheck(
team_id="XX7AN23E0Z",
bundle_id="com.akentev.app",
key_id="JSAD983ENA",
private_key=PRIVATE_KEY,
dev_environment=True
)
2 changes: 1 addition & 1 deletion tests/integration/sanic_tests/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
from sanic.response import text

from devicecheck.decorators import async_validate_device, DCSupportedAsyncFrameworks
from tests.integration.devicecheck_mock import device_check
from tests.integration.asynciodevicecheck_mock import device_check

app = Sanic(__name__)
INVALID_TOKEN_RESPONSE = text('Invalid device_token', 403)
Expand Down

0 comments on commit 45e3cc6

Please sign in to comment.