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