Bot commit: Добавлен файл token_manager.py
This commit is contained in:
parent
3e690b0e05
commit
1f33f8aeed
334
token_manager.py
Normal file
334
token_manager.py
Normal file
@ -0,0 +1,334 @@
|
|||||||
|
import os
|
||||||
|
import json
|
||||||
|
import hashlib
|
||||||
|
import base64
|
||||||
|
from datetime import datetime, timedelta, UTC
|
||||||
|
import httpx
|
||||||
|
import logging
|
||||||
|
from bson import ObjectId
|
||||||
|
from config import *
|
||||||
|
from db import mongo_client, redis_client, gridfs_map
|
||||||
|
import asyncio
|
||||||
|
from utils import fetch_bytes
|
||||||
|
import io
|
||||||
|
import uuid
|
||||||
|
from utils import encrypt, decrypt
|
||||||
|
|
||||||
|
# ---------------------- Константы OAuth ----------------------
|
||||||
|
OAUTH_URL = 'https://id.vk.com/oauth2/auth'
|
||||||
|
AUTHORIZE_URL = 'https://id.vk.ru/authorize'
|
||||||
|
SCOPES = 'offline,wall,photos,groups'
|
||||||
|
REDIRECT_URL = f'{FULL_DOMAIN}/api/v1/social/vk/token'
|
||||||
|
|
||||||
|
|
||||||
|
class TokenManager:
|
||||||
|
"""
|
||||||
|
Управление VK-токенами с шифрованием чувствительных данных.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
self.tokens_coll = mongo_client["vk"]["tokens"]
|
||||||
|
|
||||||
|
async def get_auth_url(self, user_id: str) -> str:
|
||||||
|
state = os.urandom(32).hex()
|
||||||
|
code_verifier = os.urandom(50).hex()
|
||||||
|
code_challenge = base64.urlsafe_b64encode(
|
||||||
|
hashlib.sha256(code_verifier.encode('utf-8')).digest()
|
||||||
|
).decode('utf-8').rstrip('=')
|
||||||
|
|
||||||
|
oauth_data = {"user_id": user_id, "code_verifier": code_verifier}
|
||||||
|
await redis_client.setex(f"oauth_state:{state}", 60 * 15, json.dumps(oauth_data))
|
||||||
|
|
||||||
|
url = (
|
||||||
|
f"{AUTHORIZE_URL}?"
|
||||||
|
f"response_type=code&client_id={CLIENT_ID}&scope={SCOPES}"
|
||||||
|
f"&redirect_uri={REDIRECT_URL}&state={state}"
|
||||||
|
f"&code_challenge={code_challenge}&code_challenge_method=S256"
|
||||||
|
)
|
||||||
|
return url
|
||||||
|
|
||||||
|
async def save_vk(self, contents):
|
||||||
|
grid = gridfs_map.get('social')
|
||||||
|
filename = str(uuid.uuid4())
|
||||||
|
stream = io.BytesIO(contents)
|
||||||
|
return await grid.upload_from_stream(filename, stream)
|
||||||
|
|
||||||
|
async def get_token(self, code: str, state: str, device_id: str = None) -> dict | None:
|
||||||
|
"""
|
||||||
|
Принимает `code` (выданный VK после авторизации) и `state`
|
||||||
|
(проверка, что запрос пришёл от того же клиента).
|
||||||
|
Возвращает словарь токенов или `None` в случае ошибки.
|
||||||
|
"""
|
||||||
|
# Проверка обязательных параметров
|
||||||
|
if not state or not code:
|
||||||
|
logging.error("[TokenManager] Missing code or state")
|
||||||
|
return None
|
||||||
|
|
||||||
|
# Читаем данные из Redis, ранее сохранённые в get_auth_url()
|
||||||
|
key = f"oauth_state:{state}"
|
||||||
|
oauth_data_raw = await redis_client.get(key)
|
||||||
|
|
||||||
|
if not oauth_data_raw:
|
||||||
|
logging.error(f"[TokenManager] State not found or expired: {state}")
|
||||||
|
return None
|
||||||
|
|
||||||
|
oauth_data = json.loads(oauth_data_raw)
|
||||||
|
user_id = oauth_data["user_id"]
|
||||||
|
code_verifier = oauth_data["code_verifier"]
|
||||||
|
|
||||||
|
# Выполняем запрос к VK для обмена кода на токены
|
||||||
|
try:
|
||||||
|
async with httpx.AsyncClient(timeout=15.0, verify=ctx, http2=True) as client:
|
||||||
|
response = await client.post(
|
||||||
|
OAUTH_URL,
|
||||||
|
data={
|
||||||
|
'grant_type': 'authorization_code',
|
||||||
|
'code': code,
|
||||||
|
'code_verifier': code_verifier,
|
||||||
|
'client_id': CLIENT_ID,
|
||||||
|
'device_id': device_id or "",
|
||||||
|
'redirect_uri': REDIRECT_URL,
|
||||||
|
# Генерируем новый state только для логирования/отслеживания.
|
||||||
|
'state': os.urandom(32).hex(),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
response.raise_for_status()
|
||||||
|
token_data = response.json()
|
||||||
|
access_token = token_data['access_token']
|
||||||
|
|
||||||
|
reqp = await client.post(f"{VK_API_URI}users.get", data={
|
||||||
|
"access_token": access_token,
|
||||||
|
"fields": "photo_max_orig,screen_name",
|
||||||
|
"v": "5.199",
|
||||||
|
})
|
||||||
|
vk_profile = reqp.json()['response'][0]
|
||||||
|
|
||||||
|
reqg = await client.post(f"{VK_API_URI}groups.get", data={
|
||||||
|
'access_token':access_token,
|
||||||
|
'extended':1,
|
||||||
|
'filter':'admin,editor,moder',
|
||||||
|
'v':'5.199',
|
||||||
|
})
|
||||||
|
vk_groups = reqg.json()['response']['items']
|
||||||
|
urls = [vk_profile['photo_max_orig']] + [i['photo_200'] for i in vk_groups]
|
||||||
|
results = await asyncio.gather(*(fetch_bytes(client, u) for u in urls))
|
||||||
|
results = [await self.save_vk(i) for i in results]
|
||||||
|
socials = [{
|
||||||
|
'name': f"{vk_profile['last_name']} {vk_profile['first_name']}",
|
||||||
|
'screen_name': vk_profile['screen_name'],
|
||||||
|
'id': vk_profile['id'],
|
||||||
|
'img': results[0],
|
||||||
|
'user_id': user_id,
|
||||||
|
'active': True,
|
||||||
|
'type': 'vk',
|
||||||
|
}]
|
||||||
|
for i,group in enumerate(vk_groups):
|
||||||
|
socials.append({
|
||||||
|
'name': group['name'],
|
||||||
|
'screen_name': group['screen_name'],
|
||||||
|
'id': group['id'],
|
||||||
|
'group': True,
|
||||||
|
'type': 'vk',
|
||||||
|
'active': False, # True if len(vk_groups) < 4 else False,
|
||||||
|
'img': results[i+1],
|
||||||
|
'user_id': user_id
|
||||||
|
})
|
||||||
|
for doc in socials:
|
||||||
|
await mongo_client['social']['social'].update_one({'id': doc['id'],'user_id': doc['user_id']},
|
||||||
|
{'$set': doc},upsert=True)
|
||||||
|
# Сохраняем полученные токены (Mongo + Redis)
|
||||||
|
await self.save_tokens(user_id, token_data, device_id, client_id=CLIENT_ID)
|
||||||
|
|
||||||
|
# Удаляем использованный state, чтобы его нельзя было переиспользовать.
|
||||||
|
await redis_client.delete(key)
|
||||||
|
|
||||||
|
return token_data
|
||||||
|
|
||||||
|
# Обработка ошибок HTTP (напр., неверный код, истёк скоп)
|
||||||
|
except httpx.HTTPStatusError as e:
|
||||||
|
logging.error(f"[TokenManager] Token exchange failed: {e.response.text}")
|
||||||
|
|
||||||
|
# Любые прочие исключения (сеть, тайм‑аут и т.д.)
|
||||||
|
except Exception as e:
|
||||||
|
logging.error(f"[TokenManager] Token exchange error: {e}")
|
||||||
|
|
||||||
|
# На любой ошибке гарантируем удаление state (чистка)
|
||||||
|
await redis_client.delete(key)
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
# ====================== ШИФРОВАНИЕ ======================
|
||||||
|
async def _encrypt_token_fields(self, token_data: dict) -> dict:
|
||||||
|
"""Шифрует access_token и refresh_token."""
|
||||||
|
encrypted = token_data.copy()
|
||||||
|
if "access_token" in encrypted and isinstance(encrypted["access_token"], (str, bytes)):
|
||||||
|
at = encrypted["access_token"].encode("utf-8") if isinstance(encrypted["access_token"], str) else encrypted["access_token"]
|
||||||
|
encrypted["encrypted_access_token"] = encrypt(at)
|
||||||
|
encrypted.pop("access_token", None)
|
||||||
|
if "refresh_token" in encrypted and isinstance(encrypted["refresh_token"], (str, bytes)):
|
||||||
|
rt = encrypted["refresh_token"].encode("utf-8") if isinstance(encrypted["refresh_token"], str) else encrypted["refresh_token"]
|
||||||
|
encrypted["encrypted_refresh_token"] = encrypt(rt)
|
||||||
|
encrypted.pop("refresh_token", None)
|
||||||
|
return encrypted
|
||||||
|
|
||||||
|
async def _decrypt_token_fields(self, token_doc: dict) -> dict:
|
||||||
|
"""Расшифровывает токены."""
|
||||||
|
decrypted = dict(token_doc)
|
||||||
|
if "encrypted_access_token" in decrypted:
|
||||||
|
try:
|
||||||
|
decrypted["access_token"] = decrypt(decrypted["encrypted_access_token"]).decode("utf-8")
|
||||||
|
decrypted.pop("encrypted_access_token", None)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
if "encrypted_refresh_token" in decrypted:
|
||||||
|
try:
|
||||||
|
decrypted["refresh_token"] = decrypt(decrypted["encrypted_refresh_token"]).decode("utf-8")
|
||||||
|
decrypted.pop("encrypted_refresh_token", None)
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
return decrypted
|
||||||
|
|
||||||
|
async def save_tokens(self, user_id: str, token_data: dict, device_id: str = None, client_id: int = CLIENT_ID):
|
||||||
|
expires_in = token_data.get("expires_in", 3600)
|
||||||
|
expires_at = datetime.now(UTC) + timedelta(seconds=expires_in)
|
||||||
|
|
||||||
|
encrypted_data = await self._encrypt_token_fields(token_data)
|
||||||
|
|
||||||
|
doc = encrypted_data.copy()
|
||||||
|
uid = ObjectId(user_id)
|
||||||
|
doc.update({
|
||||||
|
"user_id": uid,
|
||||||
|
"expires_at": expires_at,
|
||||||
|
"client_id": client_id or CLIENT_ID,
|
||||||
|
"device_id": device_id,
|
||||||
|
})
|
||||||
|
|
||||||
|
token_json = json.dumps({
|
||||||
|
**{k: v for k, v in doc.items() if k not in ["expires_at", "user_id"]},
|
||||||
|
"expires_at": expires_at.isoformat(),
|
||||||
|
"user_id": user_id,
|
||||||
|
})
|
||||||
|
|
||||||
|
await redis_client.setex(f"tokens:{user_id}", 60 * 60 * 24 * 60, token_json)
|
||||||
|
|
||||||
|
# Зашифрованный access_token
|
||||||
|
if "access_token" in token_data:
|
||||||
|
at_bytes = token_data["access_token"].encode("utf-8") if isinstance(token_data["access_token"], str) else token_data["access_token"]
|
||||||
|
enc_access = encrypt(at_bytes)
|
||||||
|
await redis_client.setex(f"access_token:{user_id}", expires_in - 60, enc_access)
|
||||||
|
|
||||||
|
doc.pop('user_id', None)
|
||||||
|
await self.tokens_coll.update_one(
|
||||||
|
{"user_id": uid},
|
||||||
|
{"$set": doc, "$setOnInsert": {"user_id": [uid]}},
|
||||||
|
upsert=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
async def _load_from_mongo(self, user_id):
|
||||||
|
doc = await self.tokens_coll.find_one(
|
||||||
|
{"user_id": ObjectId(user_id)},
|
||||||
|
{
|
||||||
|
"expires_at": 1,
|
||||||
|
"encrypted_access_token": 1,
|
||||||
|
"encrypted_refresh_token": 1,
|
||||||
|
"device_id": 1,
|
||||||
|
"client_id": 1,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
if doc:
|
||||||
|
return await self._decrypt_token_fields(doc)
|
||||||
|
return None
|
||||||
|
|
||||||
|
async def _cache_in_redis(self, user_id, token_doc):
|
||||||
|
token_json = json.dumps({
|
||||||
|
**{k: v for k, v in token_doc.items() if k not in ["expires_at", "user_id", "_id"]},
|
||||||
|
"expires_at": token_doc['expires_at'].isoformat(),
|
||||||
|
"user_id": user_id,
|
||||||
|
})
|
||||||
|
await redis_client.setex(f"tokens:{user_id}", 60 * 60 * 24 * 60, token_json)
|
||||||
|
|
||||||
|
async def get_access_token(self, user_id: str, need_refresh: bool = False) -> str | None:
|
||||||
|
# Короткий токен (зашифрованный)
|
||||||
|
access_token_enc = await redis_client.get(f"access_token:{user_id}")
|
||||||
|
if access_token_enc and not need_refresh:
|
||||||
|
try:
|
||||||
|
return decrypt(access_token_enc).decode("utf-8")
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
# Полный кэш
|
||||||
|
cached = await redis_client.get(f"tokens:{user_id}")
|
||||||
|
if cached:
|
||||||
|
try:
|
||||||
|
token_data = json.loads(cached)
|
||||||
|
token_data = await self._decrypt_token_fields(token_data)
|
||||||
|
except Exception:
|
||||||
|
token_data = None
|
||||||
|
else:
|
||||||
|
token_data = None
|
||||||
|
|
||||||
|
if isinstance(token_data, dict) and "error" in token_data:
|
||||||
|
token_data = None
|
||||||
|
|
||||||
|
if token_data is None:
|
||||||
|
doc = await self._load_from_mongo(user_id)
|
||||||
|
if not doc:
|
||||||
|
return None
|
||||||
|
await self._cache_in_redis(user_id, doc)
|
||||||
|
token_data = doc
|
||||||
|
|
||||||
|
# Проверка срока
|
||||||
|
if isinstance(token_data.get("expires_at"), str):
|
||||||
|
expires_at = datetime.fromisoformat(token_data["expires_at"].replace("Z", "+00:00"))
|
||||||
|
else:
|
||||||
|
expires_at = token_data.get("expires_at")
|
||||||
|
|
||||||
|
if expires_at and expires_at.astimezone(UTC) > datetime.now(UTC) + timedelta(minutes=5) and not need_refresh:
|
||||||
|
ttl = int((expires_at.astimezone(UTC) - datetime.now(UTC)).total_seconds()) - 30
|
||||||
|
enc = encrypt(token_data["access_token"].encode("utf-8"))
|
||||||
|
await redis_client.setex(f"access_token:{user_id}", ttl, enc)
|
||||||
|
return token_data["access_token"]
|
||||||
|
|
||||||
|
return await self.refresh_tokens(
|
||||||
|
user_id,
|
||||||
|
token_data.get("refresh_token"),
|
||||||
|
token_data.get("device_id"),
|
||||||
|
token_data.get('client_id', CLIENT_ID)
|
||||||
|
)
|
||||||
|
|
||||||
|
async def refresh_tokens(self, user_id: str, refresh_token: str, device_id: str, client_id: int = CLIENT_ID) -> str | None:
|
||||||
|
if not refresh_token or not device_id:
|
||||||
|
return None
|
||||||
|
try:
|
||||||
|
async with httpx.AsyncClient(timeout=15.0, verify=ctx, http2=True) as client:
|
||||||
|
resp = await client.post(
|
||||||
|
OAUTH_URL,
|
||||||
|
data={
|
||||||
|
'grant_type': 'refresh_token',
|
||||||
|
'refresh_token': refresh_token,
|
||||||
|
'client_id': client_id or CLIENT_ID,
|
||||||
|
'device_id': device_id,
|
||||||
|
'state': os.urandom(32).hex(),
|
||||||
|
},
|
||||||
|
)
|
||||||
|
resp.raise_for_status()
|
||||||
|
new_tokens = resp.json()
|
||||||
|
|
||||||
|
# В случае ошибок VK вернёт поле `error`
|
||||||
|
if 'error' in new_tokens:
|
||||||
|
logging.error(
|
||||||
|
f"[TokenManager] Refresh failed for user {user_id}: "
|
||||||
|
f"{new_tokens['error']} - {new_tokens.get('error_description', 'No description')}"
|
||||||
|
)
|
||||||
|
return None
|
||||||
|
|
||||||
|
# Сохраняем fresh‑токены
|
||||||
|
await self.save_tokens(user_id, new_tokens, device_id, client_id)
|
||||||
|
return new_tokens.get("access_token")
|
||||||
|
except Exception as e:
|
||||||
|
logging.error(f"[TokenManager] Refresh failed: {e}")
|
||||||
|
return None
|
||||||
|
|
||||||
|
async def revoke_tokens(self, user_id: str):
|
||||||
|
await self.tokens_coll.delete_one({"user_id": ObjectId(user_id)})
|
||||||
|
await redis_client.delete(f"tokens:{user_id}", f"access_token:{user_id}")
|
||||||
Loading…
Reference in New Issue
Block a user