#!/usr/bin/env python3
"""
imap_checker.py — проверка доступности IMAP по flow Thunderbird.
Использование:
    python imap_checker.py emails.txt
    python imap_checker.py emails.txt --auth   (с проверкой логина)

emails.txt — файл, по одной записи на строку:
    user@example.com
    user@example.com:password    (если нужен --auth)
"""

import argparse
import csv
import dns.resolver
import imaplib
import json
import socket
import ssl
import sys
import urllib.request
import xml.etree.ElementTree as ET
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass, field, asdict
from pathlib import Path
from typing import Optional


@dataclass
class ImapServer:
    hostname: str
    port: int
    ssl_type: str  # "SSL" | "STARTTLS" | "NONE"
    source: str    # откуда узнали: autoconfig / ispdb / mx / guess


@dataclass
class CheckResult:
    email: str
    server: Optional[str] = None
    port: Optional[int] = None
    ssl_type: Optional[str] = None
    source: Optional[str] = None
    connectable: bool = False
    auth_ok: Optional[bool] = None
    banner: Optional[str] = None
    error: Optional[str] = None
    candidates_tried: int = 0


# ─── 1. Autoconfig (Thunderbird-style) ───────────────────────────────

def fetch_autoconfig(domain: str, timeout: float = 5.0) -> list[ImapServer]:
    """Пробуем autoconfig.domain/mail/config-v1.1.xml и domain/.well-known/autoconfig/"""
    urls = [
        f"https://autoconfig.{domain}/mail/config-v1.1.xml",
        f"https://{domain}/.well-known/autoconfig/mail/config-v1.1.xml",
    ]
    servers = []
    for url in urls:
        try:
            req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
            with urllib.request.urlopen(req, timeout=timeout) as resp:
                xml_data = resp.read()
            servers.extend(_parse_autoconfig_xml(xml_data, source="autoconfig"))
            if servers:
                return servers
        except Exception:
            continue
    return servers


def fetch_ispdb(domain: str, timeout: float = 5.0) -> list[ImapServer]:
    """Mozilla ISPDB — центральная база конфигов."""
    url = f"https://autoconfig.thunderbird.net/v1.1/{domain}"
    try:
        req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            xml_data = resp.read()
        return _parse_autoconfig_xml(xml_data, source="ispdb")
    except Exception:
        return []


def _parse_autoconfig_xml(xml_data: bytes, source: str) -> list[ImapServer]:
    """Парсим XML autoconfig, вытаскиваем incomingServer type=imap."""
    servers = []
    try:
        root = ET.fromstring(xml_data)
        for inc in root.iter("incomingServer"):
            if inc.get("type") != "imap":
                continue
            hostname_el = inc.find("hostname")
            port_el = inc.find("port")
            ssl_el = inc.find("socketType")
            if hostname_el is None or port_el is None:
                continue
            hostname = hostname_el.text.strip()
            port = int(port_el.text.strip())
            ssl_type = (ssl_el.text.strip().upper() if ssl_el is not None else "SSL")
            if ssl_type == "SSL/TLS":
                ssl_type = "SSL"
            servers.append(ImapServer(hostname, port, ssl_type, source))
    except ET.ParseError:
        pass
    return servers


# ─── 2. MX-lookup ────────────────────────────────────────────────────

def discover_from_mx(domain: str) -> list[ImapServer]:
    """Достаём MX-записи, пробуем IMAP на MX-хостах."""
    servers = []
    try:
        answers = dns.resolver.resolve(domain, "MX")
        for rdata in sorted(answers, key=lambda r: r.preference):
            mx_host = str(rdata.exchange).rstrip(".")
            servers.append(ImapServer(mx_host, 993, "SSL", "mx"))
            servers.append(ImapServer(mx_host, 143, "STARTTLS", "mx"))
    except Exception:
        pass
    return servers


# ─── 3. Guess (стандартные имена) ─────────────────────────────────────

def guess_servers(domain: str) -> list[ImapServer]:
    """Перебираем типичные имена серверов, как делает Thunderbird."""
    prefixes = ["imap", "mail", "mx", "imap4", "inbox"]
    servers = []
    for prefix in prefixes:
        host = f"{prefix}.{domain}"
        servers.append(ImapServer(host, 993, "SSL", "guess"))
        servers.append(ImapServer(host, 143, "STARTTLS", "guess"))
    servers.append(ImapServer(domain, 993, "SSL", "guess"))
    servers.append(ImapServer(domain, 143, "STARTTLS", "guess"))
    return servers


# ─── 4. Проба подключения ────────────────────────────────────────────

def try_connect(srv: ImapServer, timeout: float = 8.0) -> tuple[bool, str]:
    """Пытаемся установить IMAP-соединение, возвращаем (ok, banner_or_error)."""
    try:
        if srv.ssl_type == "SSL":
            ctx = ssl.create_default_context()
            conn = imaplib.IMAP4_SSL(srv.hostname, srv.port, ssl_context=ctx, timeout=timeout)
        else:
            conn = imaplib.IMAP4(srv.hostname, srv.port, timeout=timeout)
            if srv.ssl_type == "STARTTLS":
                ctx = ssl.create_default_context()
                conn.starttls(ssl_context=ctx)

        banner = conn.welcome.decode(errors="replace") if conn.welcome else ""
        conn.logout()
        return True, banner
    except Exception as e:
        return False, str(e)


def try_auth(srv: ImapServer, user: str, password: str, timeout: float = 10.0) -> tuple[bool, str]:
    """Подключаемся + логинимся."""
    try:
        if srv.ssl_type == "SSL":
            ctx = ssl.create_default_context()
            conn = imaplib.IMAP4_SSL(srv.hostname, srv.port, ssl_context=ctx, timeout=timeout)
        else:
            conn = imaplib.IMAP4(srv.hostname, srv.port, timeout=timeout)
            if srv.ssl_type == "STARTTLS":
                ctx = ssl.create_default_context()
                conn.starttls(ssl_context=ctx)

        conn.login(user, password)
        conn.logout()
        return True, "auth_ok"
    except imaplib.IMAP4.error as e:
        return False, f"auth_fail: {e}"
    except Exception as e:
        return False, str(e)


# ─── 5. Основной flow ────────────────────────────────────────────────

def check_email(email: str, password: Optional[str] = None, do_auth: bool = False) -> CheckResult:
    """Полный цикл проверки одного адреса по flow Thunderbird."""
    result = CheckResult(email=email)

    parts = email.split("@")
    if len(parts) != 2:
        result.error = "invalid_email_format"
        return result

    domain = parts[1].lower()

    # собираем кандидатов в порядке приоритета Thunderbird
    candidates: list[ImapServer] = []
    candidates.extend(fetch_autoconfig(domain))
    candidates.extend(fetch_ispdb(domain))
    candidates.extend(discover_from_mx(domain))
    candidates.extend(guess_servers(domain))

    # дедупликация по (host, port)
    seen = set()
    unique: list[ImapServer] = []
    for c in candidates:
        key = (c.hostname.lower(), c.port)
        if key not in seen:
            seen.add(key)
            unique.append(c)

    result.candidates_tried = len(unique)

    for srv in unique:
        ok, info = try_connect(srv)
        if ok:
            result.server = srv.hostname
            result.port = srv.port
            result.ssl_type = srv.ssl_type
            result.source = srv.source
            result.connectable = True
            result.banner = info[:200]

            if do_auth and password:
                auth_ok, auth_info = try_auth(srv, email, password)
                result.auth_ok = auth_ok
                if not auth_ok:
                    result.error = auth_info

            return result

    result.error = "no_connectable_server_found"
    return result


# ─── CLI ──────────────────────────────────────────────────────────────

def parse_line(line: str) -> tuple[str, Optional[str]]:
    """Парсим строку: email или email:password."""
    line = line.strip()
    if ":" in line:
        email, password = line.split(":", 1)
        return email.strip(), password.strip()
    return line, None


def main():
    parser = argparse.ArgumentParser(description="IMAP availability checker (Thunderbird flow)")
    parser.add_argument("input", help="Файл с email-адресами (по одному на строку, опционально :password)")
    parser.add_argument("--auth", action="store_true", help="Проверять аутентификацию (нужны пароли в файле)")
    parser.add_argument("--threads", type=int, default=10, help="Количество потоков (по умолчанию 10)")
    parser.add_argument("--output", "-o", help="Файл для результатов (JSON). По умолчанию — stdout")
    parser.add_argument("--csv", dest="csv_out", help="Файл для результатов в CSV")
    args = parser.parse_args()

    lines = Path(args.input).read_text(encoding="utf-8").strip().splitlines()
    entries = [parse_line(l) for l in lines if l.strip() and not l.strip().startswith("#")]

    results: list[CheckResult] = []

    print(f"[*] Проверяю {len(entries)} адресов, потоков: {args.threads}", file=sys.stderr)

    with ThreadPoolExecutor(max_workers=args.threads) as pool:
        futures = {
            pool.submit(check_email, email, password, args.auth): email
            for email, password in entries
        }
        for i, future in enumerate(as_completed(futures), 1):
            res = future.result()
            results.append(res)

            status = "OK" if res.connectable else "FAIL"
            auth_str = ""
            if res.auth_ok is not None:
                auth_str = f" auth={'OK' if res.auth_ok else 'FAIL'}"

            server_str = f" -> {res.server}:{res.port} ({res.ssl_type}, {res.source})" if res.server else ""
            print(f"[{i}/{len(entries)}] {status} {res.email}{server_str}{auth_str}", file=sys.stderr)

    # вывод
    data = [asdict(r) for r in results]

    if args.output:
        Path(args.output).write_text(json.dumps(data, indent=2, ensure_ascii=False), encoding="utf-8")
        print(f"[+] JSON сохранён: {args.output}", file=sys.stderr)
    elif not args.csv_out:
        print(json.dumps(data, indent=2, ensure_ascii=False))

    if args.csv_out:
        with open(args.csv_out, "w", newline="", encoding="utf-8") as f:
            writer = csv.DictWriter(f, fieldnames=list(data[0].keys()) if data else [])
            writer.writeheader()
            writer.writerows(data)
        print(f"[+] CSV сохранён: {args.csv_out}", file=sys.stderr)

    # итоги
    total = len(results)
    ok = sum(1 for r in results if r.connectable)
    auth_ok = sum(1 for r in results if r.auth_ok is True)
    print(f"\n[=] Итого: {ok}/{total} доступны по IMAP", file=sys.stderr)
    if args.auth:
        print(f"[=] Аутентификация: {auth_ok}/{total} успешно", file=sys.stderr)


if __name__ == "__main__":
    main()
