334 lines
14 KiB
Python
334 lines
14 KiB
Python
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}") |