1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
|
"""Support for August devices."""
from __future__ import annotations
import time
from collections import defaultdict
from ..exceptions import RateLimited
RATE_LIMIT_WAKEUP_INTERVAL = 60 * 26
_NEVER_TIME = -RATE_LIMIT_WAKEUP_INTERVAL
class RateLimitCheck:
"""The rate limit is checked locally here to avoid getting blocked.
This is a basic rate limit checker that will check the rate limit
locally to avoid getting blocked.
If rate limiting is not checked locally there is a risk that the
client will get permanently blocked by the server.
"""
def __init__(self) -> None:
"""Initialize the rate limit checker."""
self._client_wakeups: defaultdict[str, float] = defaultdict(lambda: _NEVER_TIME)
async def check_rate_limit(self, token: str) -> None:
"""Check if the client is rate limited."""
now = time.monotonic()
last_time = self._client_wakeups[token]
next_allowed = last_time + RATE_LIMIT_WAKEUP_INTERVAL
if next_allowed > now:
min_until_next_allowed = int((next_allowed - now) / 60)
raise RateLimited(
f"Rate limited, try again in {min_until_next_allowed} minutes",
next_allowed,
)
async def register_wakeup(self, token: str) -> None:
"""Register a wakeup for the client."""
self._client_wakeups[token] = time.monotonic()
_RateLimitChecker = RateLimitCheck()
|
