TDATA-Checker/modules/checker.py

import logging

from dataclasses import dataclass
from pathlib import Path
from phonenumbers import parse, geocoder

from typing import List, Optional

from telethon.types import User, Chat, Channel
from telethon.tl.functions.payments import GetStarsStatusRequest
from telethon.tl.functions.contacts import GetContactsRequest

from modules.tdata import TDATA

logger = logging.getLogger(__name__)


class SessionInvalid(Exception):
    pass


# Unified Result class
@dataclass
class Result:
    me: Optional[User] = None
    stars: Optional[int] = None
    country: Optional[str] = None

    contacts: Optional[List[User]] = None
    chats: Optional[List[Chat]] = None
    channels: Optional[List[Channel]] = None
    admin_channels: Optional[List[Channel]] = None


class Checker:
    def __init__(self, path: Path, temp_path: Path):
        self.path = path
        self.temp_path = temp_path
        self.client = None
        logger.debug(f"Создание экземпляра Checker для пути: {path}")
        try:
            self.desktop = TDATA(path=self.path, temp_path=self.temp_path)
            logger.debug("TDATA успешно инициализирована")
        except Exception as e:
            logger.error(f"Ошибка инициализации TDATA: {e}")
            self.desktop = None

    async def __aenter__(self):
        logger.debug("Вход в контекстный менеджер Checker")
        await self.__ensure_client()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        logger.debug("Выход из контекстного менеджера Checker")
        if self.client:
            try:
                await self.client.disconnect()
                logger.debug("Telethon клиент отключен")
            except Exception as e:
                logger.warning(f"Ошибка при отключении клиента: {e}")

    async def __ensure_client(self):
        if self.client is None:
            if self.desktop is None:
                logger.error("Попытка инициализации клиента с неинициализированной TDATA")
                raise Exception("TDATA не инициализирована")

            logger.debug("Конвертация TDATA в Telethon клиент")
            self.client = await self.desktop.convert()
            logger.debug("Telethon клиент успешно создан")

    async def check_stars(self, me: User) -> int:
        logger.debug("Запрос баланса звезд")
        try:
            stars_status = await self.client(GetStarsStatusRequest(peer=me))
            stars = stars_status.balance.amount
            logger.debug(f"Баланс звезд: {stars}")
            return stars

        except Exception as e:
            logger.warning(f"Не удалось получить баланс звезд: {e}")
            return 0

    async def check_country(self, me: User) -> str:
        logger.debug(f"Определение страны для номера +{me.phone}")
        try:
            phone_parse = parse(str(f"+{me.phone}"))
            country = geocoder.description_for_number(phone_parse, "ru")
            logger.debug(f"Страна определена: {country}")
            return country

        except Exception as e:
            logger.warning(f"Не удалось определить страну: {e}")
            return "N/A"

    async def check_contacts(self) -> List[User]:
        logger.debug("Запрос списка контактов")
        try:
            result = await self.client(GetContactsRequest(hash=0))
            contacts = result.users
            logger.debug(f"Получено контактов: {len(contacts)}")
            return contacts

        except Exception as e:
            logger.warning(f"Не удалось получить контакты: {e}")
            return []

    async def check_chats(self) -> List[Chat]:
        logger.debug("Запрос всех диалогов")
        try:
            chats = []
            async for dialog in self.client.iter_dialogs():
                chats.append(dialog.entity)

            logger.debug(f"Получено диалогов: {len(chats)}")
            return chats

        except Exception as e:
            logger.warning(f"Не удалось получить чаты: {e}")
            return []

    async def check_channels(self, chats: List[Chat]) -> List[Channel]:
        logger.debug(f"Фильтрация каналов из {len(chats)} чатов")
        try:
            channels = [chat for chat in chats if isinstance(chat, Channel)]
            logger.debug(f"Найдено каналов: {len(channels)}")
            return channels

        except Exception as e:
            logger.warning(f"Не удалось отфильтровать каналы: {e}")
            return []

    async def check_admin_channels(self, channels: List[Channel]) -> List[Channel]:
        logger.debug("Проверка прав администратора в каналах")
        try:
            admin_channels = []
            checked = 0

            for channel in channels:
                checked += 1
                try:
                    permissions = await self.client.get_permissions(channel, "me")
                    if permissions.is_admin or permissions.is_creator:
                        admin_channels.append(channel)
                        logger.debug(f"Найден админ канал: {channel.title}")
                except Exception as e:
                    logger.debug(f"Не удалось проверить права для канала: {e}")

            logger.debug(f"Проверено каналов: {checked}, найдено админских: {len(admin_channels)}")
            return admin_channels

        except Exception as e:
            logger.warning(f"Не удалось проверить админские права: {e}")
            return []

    async def check_all(self) -> Result:
        logger.debug("Начало полной проверки аккаунта")
        try:
            me = await self.client.get_me()
            if not me:
                logger.error("Не удалось получить информацию о пользователе")
                raise SessionInvalid()

            logger.debug(f"Получен пользователь: {me.first_name} (+{me.phone})")

            stars = await self.check_stars(me=me)
            country = await self.check_country(me=me)
            contacts = await self.check_contacts()
            chats = await self.check_chats()
            channels = await self.check_channels(chats=chats) if chats else []
            admin_channels = await self.check_admin_channels(channels=channels) if channels else []

            logger.debug("Полная проверка завершена успешно")
            return Result(
                me=me,
                stars=stars,
                country=country,
                contacts=contacts,
                chats=chats,
                channels=channels,
                admin_channels=admin_channels,
            )

        except Exception as e:
            logger.error(f"Критическая ошибка при проверке: {e}")
            raise SessionInvalid()