bot.py

import os
import asyncio
from datetime import datetime
from telethon import TelegramClient, events
from telethon.tl.types import (
    MessageMediaPhoto, MessageMediaDocument,
    MessageMediaContact, MessageMediaGeo, MessageMediaPoll
)
from telethon.errors import FloodWaitError
import config
from database import Database
from logger import setup_logger

log = setup_logger('bot')


class TelegramLogger:
    def __init__(self):
        self.client = TelegramClient(
            config.SESSION_NAME,
            config.API_ID,
            config.API_HASH
        )
        self.db = Database()
        log.info("TelegramLogger initialized")
    
    async def start(self):
        await self.client.start()
        me = await self.client.get_me()
        log.info(f"Logged in as: {me.first_name} (@{me.username})")
        log.info("Starting monitoring all chats...")
        
        self.register_handlers()
        await self.client.run_until_disconnected()
    
    def register_handlers(self):
        @self.client.on(events.NewMessage)
        async def handle_new_message(event):
            await self.process_message(event.message, is_new=True)
        
        @self.client.on(events.MessageEdited)
        async def handle_edited_message(event):
            await self.process_message_edit(event.message)
        
        @self.client.on(events.MessageDeleted)
        async def handle_deleted_message(event):
            await self.process_deleted_messages(event)
    
    async def process_message(self, message, is_new=True, is_edited=False):
        try:
            chat = await message.get_chat()
            chat_id = message.chat_id
            
            if hasattr(chat, 'title'):
                chat_title = chat.title
                chat_type = 'group' if message.is_group else 'channel'
            else:
                chat_title = f"{chat.first_name or ''} {chat.last_name or ''}".strip()
                chat_type = 'private'
            
            chat_username = getattr(chat, 'username', None)
            
            self.db.save_chat(chat_id, chat_type, chat_title, chat_username)
            
            sender = await message.get_sender()
            if sender:
                sender_id = sender.id
                if hasattr(sender, 'first_name'):
                    sender_name = f"{sender.first_name or ''} {sender.last_name or ''}".strip()
                else:
                    sender_name = sender.title if hasattr(sender, 'title') else 'Unknown'
            else:
                sender_id = 0
                sender_name = 'Unknown'
            
            media_type = None
            media_path = None
            
            if message.media:
                media_type = self.get_media_type(message.media)
                media_path = await self.download_media(message, chat_id)
            
            message_data = {
                'message_id': message.id,
                'chat_id': chat_id,
                'sender_id': sender_id,
                'sender_name': sender_name,
                'text': message.text or '',
                'date': message.date,
                'media_type': media_type,
                'media_path': media_path,
                'reply_to_msg_id': message.reply_to_msg_id,
                'raw_data': {
                    'is_reply': message.is_reply,
                    'mentioned': message.mentioned,
                    'out': message.out,
                }
            }
            
            self.db.save_message(message_data, mark_edited=is_edited)
            
            if is_new:
                status = "📨"
                if media_type:
                    status = f"📎{media_type}"
                log.info(f"{status} [{chat_title}] {sender_name}: {message.text[:50] if message.text else '(media)'}")
        
        except Exception as e:
            log.error(f"Error processing message: {e}")
    
    async def process_message_edit(self, message):
        try:
            messages = self.db.get_messages(message.chat_id, limit=1000)
            old_message = next((m for m in messages if m['message_id'] == message.id), None)
            
            if old_message:
                old_text = old_message['text'] or ''
                new_text = message.text or ''
                
                old_media_info = f" [media: {old_message.get('media_type')}]" if old_message.get('media_type') else ""
                new_media_type = self.get_media_type(message.media) if message.media else None
                new_media_info = f" [media: {new_media_type}]" if new_media_type else ""
                
                old_full = old_text + old_media_info
                new_full = new_text + new_media_info
                
                if old_full != new_full or old_text != new_text:
                    self.db.save_message_edit(message.id, message.chat_id, old_full, new_full)
                    
                    chat = await message.get_chat()
                    chat_title = getattr(chat, 'title', 'Private')
                    log.info(f"✏️ [{chat_title}] Message #{message.id} edited: '{old_text[:30]}' → '{new_text[:30]}'")
                else:
                    log.warning(f"⚠️ Edit event for #{message.id}, but no changes detected")
            else:
                log.warning(f"⚠️ Message #{message.id} in chat {message.chat_id} edited, but not found in DB")
            
            await self.process_message(message, is_new=False, is_edited=True)
        
        except Exception as e:
            log.error(f"Error processing message edit: {e}", exc_info=True)
    
    async def process_deleted_messages(self, event):
        try:
            chat_id = None
            
            if hasattr(event, 'chat_id') and event.chat_id:
                chat_id = event.chat_id
            if hasattr(event, 'original_update'):
                update = event.original_update
                if hasattr(update, 'channel_id'):
                    channel_id = update.channel_id
                    chat_id = int(f"-100{channel_id}")
                elif hasattr(update, 'peer'):
                    if hasattr(update.peer, 'channel_id'):
                        channel_id = update.peer.channel_id
                        chat_id = int(f"-100{channel_id}")
                    elif hasattr(update.peer, 'chat_id'):
                        chat_id = -update.peer.chat_id
                    elif hasattr(update.peer, 'user_id'):
                        chat_id = update.peer.user_id
            
            for msg_id in event.deleted_ids:
                if not chat_id:
                    chat_id = self.db.find_message_chat_id(msg_id)
                    if not chat_id:
                        log.warning(f"⚠️ Could not find chat_id for deleted message #{msg_id}")
                        continue
                
                success = self.db.mark_message_deleted(msg_id, chat_id)
                if success:
                    log.info(f"🗑️ Message #{msg_id} deleted in chat {chat_id}")
                else:
                    log.warning(f"⚠️ Message #{msg_id} not found in DB")
        
        except Exception as e:
            log.error(f"Error processing message deletion: {e}", exc_info=True)
    
    def get_media_type(self, media) -> str:
        if isinstance(media, MessageMediaPhoto):
            return 'photo'
        elif isinstance(media, MessageMediaDocument):
            if hasattr(media.document, 'mime_type'):
                mime = media.document.mime_type.lower()
                if 'video' in mime:
                    return 'video'
                elif 'audio' in mime:
                    return 'audio'
                elif 'voice' in mime or 'ogg' in mime:
                    return 'voice'
            if hasattr(media.document, 'attributes'):
                for attr in media.document.attributes:
                    attr_type = type(attr).__name__
                    if 'Video' in attr_type:
                        return 'video'
                    elif 'Audio' in attr_type:
                        return 'audio'
                    elif 'Voice' in attr_type:
                        return 'voice'
            return 'document'
        elif isinstance(media, MessageMediaContact):
            return 'contact'
        elif isinstance(media, MessageMediaGeo):
            return 'location'
        elif isinstance(media, MessageMediaPoll):
            return 'poll'
        else:
            return 'other'
    
    async def download_media(self, message, chat_id) -> str:
        try:
            chat_folder = os.path.join(config.MEDIA_FOLDER, str(chat_id))
            os.makedirs(chat_folder, exist_ok=True)
            
            timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
            filename = f"{timestamp}_{message.id}"
            
            file_path = await self.client.download_media(
                message.media,
                file=os.path.join(chat_folder, filename)
            )
            
            if file_path:
                return os.path.relpath(file_path, config.MEDIA_FOLDER)
            
        except FloodWaitError as e:
            log.warning(f"⏳ Flood wait: need to wait {e.seconds} seconds")
        except Exception as e:
            log.error(f"Error downloading media: {e}")
        
        return None
    
    async def stop(self):
        await self.client.disconnect()
        log.info("✓ UserBot stopped")


async def main():
    logger = TelegramLogger()
    
    try:
        await logger.start()
    except KeyboardInterrupt:
        log.warning("\n⚠️ Shutdown signal received...")
        await logger.stop()
    except Exception as e:
        log.error(f"Critical error: {e}", exc_info=True)
        await logger.stop()


if __name__ == '__main__':
    log.info("=" * 50)
    log.info("  Telegram Logger - UserBot")
    log.info("=" * 50)
    asyncio.run(main())