36 lines
1.0 KiB
Python
36 lines
1.0 KiB
Python
|
import random
|
||
|
|
import string
|
||
|
|
import time
|
||
|
|
from collections import defaultdict
|
||
|
|
|
||
|
|
|
||
|
|
class WeComPairing:
|
||
|
|
CODE_TTL = 600
|
||
|
|
RATE_LIMIT_WINDOW = 60
|
||
|
|
RATE_LIMIT_MAX = 5
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
self._codes: dict[str, tuple[str, float]] = {}
|
||
|
|
self._rate_limit: dict[str, list[float]] = defaultdict(list)
|
||
|
|
|
||
|
|
def generate_code(self, peer_id: str) -> str | None:
|
||
|
|
now = time.time()
|
||
|
|
timestamps = self._rate_limit[peer_id]
|
||
|
|
timestamps[:] = [t for t in timestamps if now - t < self.RATE_LIMIT_WINDOW]
|
||
|
|
|
||
|
|
if len(timestamps) >= self.RATE_LIMIT_MAX:
|
||
|
|
return None
|
||
|
|
|
||
|
|
timestamps.append(now)
|
||
|
|
code = "".join(random.choices(string.digits, k=6))
|
||
|
|
self._codes[peer_id] = (code, now + self.CODE_TTL)
|
||
|
|
return code
|
||
|
|
|
||
|
|
def verify(self, peer_id: str, code: str) -> bool:
|
||
|
|
entry = self._codes.get(peer_id)
|
||
|
|
if entry is None:
|
||
|
|
return False
|
||
|
|
stored_code, expires_at = entry
|
||
|
|
del self._codes[peer_id]
|
||
|
|
return stored_code == code and time.time() < expires_at
|