diff --git a/token_manager.py b/token_manager.py new file mode 100644 index 0000000..3742e07 --- /dev/null +++ b/token_manager.py @@ -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}") \ No newline at end of file