processor.py - master - Eugeny/MailReceiptScripts_Python

processor.py
150 строк6 KB
Eugeny
upload files
21 дек 2025, 12:51
21 дек 2025, 12:51a6a057b
import imaplib
import email
import ssl
from pathlib import Path
from typing import Optional
from config import Config
from logger import logger
from spam_filter import is_spam
from utils import decode_str, sanitize_filename, is_allowed_extension

def save_attachment(part: email.message.Message, folder: Path) -> bool:
    try:
        filename = part.get_filename()
        if not filename:
            return False
        filename = decode_str(filename)
        if not filename or not is_allowed_extension(filename):
            return False

        filename = sanitize_filename(filename)
        filepath = folder / filename

        counter = 1
        while filepath.exists():
            filepath = filepath.with_name(f"{filepath.stem}_{counter}{filepath.suffix}")
            counter += 1

        with open(filepath, "wb") as f:
            f.write(part.get_payload(decode=True))
        
        logger.info(f"Сохранено вложение: {filepath}")
        return True
    except Exception as e:
        logger.error(f"Ошибка сохранения вложения: {e}")
        return False

def get_email_body(msg: email.message.Message) -> str:
    body = []
    if msg.is_multipart():
        for part in msg.walk():
            if "attachment" in str(part.get("Content-Disposition", "")).lower():
                continue
            if part.get_content_type() == "text/plain":
                payload = part.get_payload(decode=True)
                if payload:
                    body.append(payload.decode('utf-8', errors='replace'))
    else:
        payload = msg.get_payload(decode=True)
        if payload:
            body.append(payload.decode('utf-8', errors='replace'))
    return "\n".join(body)

def connect_to_mail(email_addr: str, password: str) -> Optional[imaplib.IMAP4_SSL]:
    try:
        context = ssl.create_default_context()
        mail = imaplib.IMAP4_SSL(Config.IMAP_SERVER, Config.IMAP_PORT, ssl_context=context)
        mail.login(email_addr, password)
        logger.info(f"Подключено: {email_addr}")
        return mail
    except Exception as e:
        logger.error(f"Ошибка подключения {email_addr}: {e}")
    return None

def process_email(
    mail: imaplib.IMAP4_SSL, 
    email_id: bytes, 
    download_folder: Path
) -> bool:
    """
    Обрабатывает одно письмо.
    Помечает как прочитанное только если были сохранены вложения.
    """
    try:
        # BODY.PEEK[] загружает ПОЛНОЕ сообщение БЕЗ пометки "прочитано". Это важно, так как не каждое письмо имеет вложения! 
        status, msg_data = mail.fetch(email_id, "(BODY.PEEK[])")
        if status != 'OK':
            logger.warning(f"Не удалось загрузить письмо {email_id.decode()}")
            return False

        email_raw = msg_data[0][1]
        mail_message = email.message_from_bytes(email_raw)

        subject = decode_str(mail_message["Subject"])
        from_addr = decode_str(mail_message["From"]) or "Unknown"
        body = get_email_body(mail_message)

        # Проверка на спам
        spam_result, spam_score = is_spam(email_raw, subject, body)
        if spam_result:
            score_info = f" (score: {spam_score})" if spam_score is not None else ""
            logger.info(f"[СПАМ{score_info}] От: {from_addr}, Тема: {subject}")
            return True

        # Сохраняем вложения
        attachments_saved = False
        for part in mail_message.walk():
            content_disposition = str(part.get("Content-Disposition", "")).lower()
            if "attachment" in content_disposition:
                if save_attachment(part, download_folder):
                    attachments_saved = True

        # Помечаем как прочитанное ТОЛЬКО если сохранены вложения
        if attachments_saved:
            mail.store(email_id, "+FLAGS", "\\Seen")
            logger.info(f"[OK] Вложения сохранены, От: {from_addr}, Тема: {subject}")
        else:
            logger.info(f"[ПРОПУСК] Нет вложений, От: {from_addr}, Тема: {subject}")

        return True

    except Exception as e:
        logger.error(f"Ошибка обработки письма {email_id}: {e}")
        return False


def process_account(account: dict) -> None:
    email_addr      = account['email']
    download_folder = account['download_folder']
    
    try:
        download_folder.mkdir(exist_ok=True)
        mail = connect_to_mail(email_addr, account['password'])
        if not mail:
            return

        try:
            mail.select("INBOX")
            status, messages = mail.search(None, "UNSEEN")
            if status != 'OK':
                return

            email_ids = messages[0].split()
            if not email_ids:
                logger.info(f"Нет новых писем для {email_addr}")
                return

            logger.info(f"Писем для обработки: {len(email_ids)} ({email_addr})")
            
            for e_id in email_ids:
                process_email(mail, e_id, download_folder)

        finally:
            try:
                mail.close()
                mail.logout()
            except:
                pass

    except Exception as e:
        logger.error(f"Ошибка аккаунта {email_addr}: {e}")