#!/usr/bin/env python3

import argparse
import concurrent.futures
import json
import os
import pwd
import re
import subprocess
import sys
import tempfile
import time
import urllib.request

from collections import defaultdict
from pathlib import Path

import dns.resolver


VERSION = "6.5"

SOURCES = [
    "https://raw.githubusercontent.com/mehrtat/disposable-email-domain/main/domains.txt",
    "https://raw.githubusercontent.com/FFraud-com/disposable-email-domains/main/disposable-email-domains.txt",
    "https://disposable.github.io/disposable-email-domains/domains.txt",
]

DOMAIN_RE = re.compile(
    r"^(?=.{1,253}$)"
    r"(?:[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?\.)+"
    r"[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?$"
)


def log(message=""):
    print(message, flush=True)


def normalize_domain(value):
    if not value:
        return None

    value = value.strip().lower().rstrip(".")

    if not value or value.startswith("#"):
        return None

    if value.startswith("@"):
        value = value[1:]

    if value.startswith("*."):
        value = value[2:]

    try:
        value = value.encode("idna").decode("ascii")
    except Exception:
        return None

    if not DOMAIN_RE.match(value):
        return None

    return value


def read_domain_file(filename):
    path = Path(filename)

    if not path.exists():
        return set()

    result = set()

    for line in path.read_text(
        encoding="utf-8",
        errors="ignore",
    ).splitlines():

        domain = normalize_domain(line)

        if domain:
            result.add(domain)

    return result


def download_source(url, timeout=120):
    request = urllib.request.Request(
        url,
        headers={
            "User-Agent":
                f"mastodon-disposable-email-blocker/{VERSION}"
        },
    )

    with urllib.request.urlopen(
        request,
        timeout=timeout,
    ) as response:

        content = response.read().decode(
            "utf-8",
            errors="ignore",
        )

    result = set()

    for line in content.splitlines():
        domain = normalize_domain(line)

        if domain:
            result.add(domain)

    return result


def chunks(values, size):
    values = list(values)

    for offset in range(
        0,
        len(values),
        size,
    ):
        yield values[
            offset:
            offset + size
        ]


def save_json_atomic(filename, value):
    filename = Path(filename)

    filename.parent.mkdir(
        parents=True,
        exist_ok=True,
    )

    fd, temporary = tempfile.mkstemp(
        prefix=filename.name + ".",
        dir=str(filename.parent),
    )

    try:
        with os.fdopen(
            fd,
            "w",
            encoding="utf-8",
        ) as handle:

            json.dump(
                value,
                handle,
                indent=2,
                sort_keys=True,
            )

            handle.flush()
            os.fsync(handle.fileno())

        os.replace(
            temporary,
            filename,
        )

    finally:
        if os.path.exists(temporary):
            os.unlink(temporary)


def mastodon_owner(cwd):
    path = Path(cwd).resolve()
    stat = path.stat()
    account = pwd.getpwuid(stat.st_uid)

    return {
        "uid": stat.st_uid,
        "gid": stat.st_gid,
        "name": account.pw_name,
        "home": Path(account.pw_dir),
    }


def command_environment(cwd):
    env = os.environ.copy()
    owner = mastodon_owner(cwd)
    home = owner["home"]

    candidates = [
        home / ".rbenv" / "shims",
        home / ".rbenv" / "bin",
        home / ".asdf" / "shims",
        home / ".asdf" / "bin",
        home / ".local" / "share" / "mise" / "shims",
        home / ".local" / "bin",
    ]

    path_parts = []

    for candidate in candidates:
        if candidate.is_dir():
            path_parts.append(str(candidate))

    current_path = env.get("PATH", "")

    if current_path:
        path_parts.append(current_path)

    env["PATH"] = os.pathsep.join(path_parts)
    env["HOME"] = str(home)
    env["USER"] = owner["name"]
    env["LOGNAME"] = owner["name"]
    env["RAILS_ENV"] = "production"

    rbenv_root = home / ".rbenv"

    if rbenv_root.is_dir():
        env["RBENV_ROOT"] = str(rbenv_root)

    return env


def demote_to_mastodon_owner(cwd):
    owner = mastodon_owner(cwd)

    if os.geteuid() != 0:
        return None

    if owner["uid"] == 0:
        return None

    def demote():
        os.initgroups(owner["name"], owner["gid"])
        os.setgid(owner["gid"])
        os.setuid(owner["uid"])

    return demote


def run_command(
    command,
    cwd,
    check=True,
    extra_env=None,
):
    env = command_environment(cwd)

    if extra_env:
        env.update(extra_env)

    process = subprocess.run(
        [str(item) for item in command],
        cwd=str(cwd),
        env=env,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        preexec_fn=demote_to_mastodon_owner(cwd),
    )

    if check and process.returncode != 0:
        if process.stdout:
            sys.stderr.write(
                process.stdout
            )

        if process.stderr:
            sys.stderr.write(
                process.stderr
            )

        raise RuntimeError(
            "Comando fallito: "
            + " ".join(
                str(item)
                for item in command
            )
        )

    return process

def get_mastodon_blocks(mastodon_dir):
    ruby = r'''
rows = EmailDomainBlock
  .order(:id)
  .pluck(
    :id,
    :domain,
    :allow_with_approval,
    :parent_id
  )

STDOUT.write(
  "__JSON__" + JSON.generate(rows)
)
'''

    process = run_command(
        [
            "bin/rails",
            "runner",
            ruby,
        ],
        cwd=mastodon_dir,
    )

    marker = "__JSON__"

    position = process.stdout.rfind(
        marker
    )

    if position < 0:
        raise RuntimeError(
            "Rails è partito ma non trovo "
            "il JSON di EmailDomainBlock"
        )

    return json.loads(
        process.stdout[
            position + len(marker):
        ]
    )


def force_everything_to_deny(
    mastodon_dir,
):
    ruby = r'''
count = EmailDomainBlock
  .where(allow_with_approval: true)
  .update_all(
    allow_with_approval: false,
    updated_at: Time.current
  )

STDOUT.write(
  "__COUNT__#{count}"
)
'''

    process = run_command(
        [
            "bin/rails",
            "runner",
            ruby,
        ],
        cwd=mastodon_dir,
    )

    marker = "__COUNT__"

    position = process.stdout.rfind(
        marker
    )

    if position < 0:
        raise RuntimeError(
            "Errore conversione "
            "approval -> deny"
        )

    return int(
        process.stdout[
            position + len(marker):
        ].strip()
    )


def tootctl_add(
    mastodon_dir,
    domains,
    dry_run=False,
):
    if not domains:
        return

    if dry_run:
        log(
            f"[DRY-RUN] aggiungerei "
            f"{len(domains):,} domini"
        )

        return

    process = run_command(
        [
            "bin/tootctl",
            "email-domain-blocks",
            "add",
            *domains,
        ],
        cwd=mastodon_dir,
    )

    if process.stdout:
        for line in (
            process.stdout.splitlines()
        ):
            if line.strip():
                log(
                    "    tootctl: "
                    + line.strip()
                )


def cleanup_blocked_pending(
    mastodon_dir,
    dry_run=False,
):

    ruby = r'''dry_run = ENV["BLOCKER_DRY_RUN"] == "1"

checked = 0
matched = 0
matched_email = 0
matched_username = 0
deleted = 0
errors = []

User.pending.find_each(batch_size: 100) do |user|
  next if user.account.nil?

  checked += 1

  begin
    blocked_email =
      user.email.present? &&
      EmailDomainBlock.block?(user.email)

    blocked_username =
      user.account.username.to_s.match?(/\\Ausr_/i)

    next unless blocked_email || blocked_username

    matched += 1
    matched_email += 1 if blocked_email
    matched_username += 1 if blocked_username

    reasons = []
    reasons << "email" if blocked_email
    reasons << "username" if blocked_username

    if dry_run
      puts "__PENDING__#{user.id}\\t#{user.account.username}\\t#{reasons.join(",")}"
      next
    end

    DeleteAccountService.new.call(
      user.account,
      reserve_email: false,
      reserve_username: false
    )

    deleted += 1

  rescue => e
    errors << {
      id: user.id,
      username: user.account&.username,
      error: "#{e.class}: #{e.message}"
    }
  end
end

STDOUT.write(
  "__CLEANUP__" + JSON.generate({
    checked: checked,
    matched: matched,
    matched_email: matched_email,
    matched_username: matched_username,
    deleted: deleted,
    errors: errors
  })
)
'''

    process = run_command(
        [
            "bin/rails",
            "runner",
            ruby,
        ],
        cwd=mastodon_dir,
        extra_env={
            "BLOCKER_DRY_RUN":
                "1" if dry_run else "0"
        },
    )

    if dry_run:
        for line in (
            process.stdout.splitlines()
        ):
            if line.startswith(
                "__PENDING__"
            ):
                data = line[
                    len("__PENDING__"):
                ]

                log(
                    "[pending] DRY-RUN: "
                    + data
                )

    marker = "__CLEANUP__"

    position = process.stdout.rfind(
        marker
    )

    if position < 0:
        raise RuntimeError(
            "Risultato cleanup pending "
            "non trovato"
        )

    return json.loads(
        process.stdout[
            position + len(marker):
        ]
    )

def resolve_mx(
    domain,
    timeout,
):
    resolver = dns.resolver.Resolver(
        configure=True
    )

    resolver.timeout = timeout
    resolver.lifetime = timeout

    try:
        answers = resolver.resolve(
            domain,
            "MX",
        )

    except Exception:
        return domain, []

    result = set()

    for answer in answers:
        mx = normalize_domain(
            str(answer.exchange)
        )

        if mx:
            result.add(mx)

    return domain, sorted(result)


def main():
    parser = argparse.ArgumentParser(
        description=(
            "Mastodon disposable email blocker. "
            "Disposable domains are hard-blocked; "
            "MX records are report-only; "
            "pending registrations using blocked "
            "email domains or usr_ usernames are removed."
        )
    )

    parser.add_argument(
        "--mastodon-dir",
        default="/home/mastodon/live",
    )

    parser.add_argument(
        "--allowlist",
        default=(
            "/etc/mastodon-email-blocker/"
            "allowlist.txt"
        ),
    )

    parser.add_argument(
        "--state",
        default=(
            "/var/lib/mastodon-email-blocker/"
            "state.json"
        ),
    )

    parser.add_argument(
        "--report",
        default=(
            "/var/lib/mastodon-email-blocker/"
            "last-report.json"
        ),
    )

    parser.add_argument(
        "--batch-size",
        type=int,
        default=250,
    )

    parser.add_argument(
        "--mx-workers",
        type=int,
        default=40,
    )

    parser.add_argument(
        "--mx-timeout",
        type=float,
        default=4.0,
    )

    parser.add_argument(
        "--mx-mode",
        choices=[
            "report",
            "off",
        ],
        default="off",
    )

    parser.add_argument(
        "--dry-run",
        action="store_true",
    )

    parser.add_argument(
        "--no-force-deny",
        action="store_true",
    )

    parser.add_argument(
        "--no-cleanup-pending",
        action="store_true",
        help=(
            "Do not remove pending registrations "
            "matched by email block or usr_ username"
        ),
    )

    args = parser.parse_args()

    mastodon_dir = Path(
        args.mastodon_dir
    )

    if not (
        mastodon_dir / "bin/rails"
    ).exists():
        raise SystemExit(
            "bin/rails non trovato in "
            + str(mastodon_dir)
        )

    if not (
        mastodon_dir / "bin/tootctl"
    ).exists():
        raise SystemExit(
            "bin/tootctl non trovato in "
            + str(mastodon_dir)
        )

    allowlist = read_domain_file(
        args.allowlist
    )

    log()
    log(
        "============================================"
    )
    log(
        f" Mastodon Disposable Email Blocker v{VERSION}"
    )
    log(
        "============================================"
    )
    log()

    log(
        f"[config] Mastodon: "
        f"{mastodon_dir}"
    )

    owner = mastodon_owner(
        mastodon_dir
    )

    log(
        f"[config] Utente Mastodon: "
        f"{owner['name']} "
        f"(uid={owner['uid']})"
    )

    ruby_version = run_command(
        ["ruby", "-v"],
        cwd=mastodon_dir,
    ).stdout.strip()

    log(
        f"[config] Ruby: {ruby_version}"
    )

    log(
        f"[config] allowlist: "
        f"{len(allowlist):,}"
    )

    log(
        f"[config] MX mode: "
        f"{args.mx_mode}"
    )

    log(
        "[config] MX autoblock: DISABILITATO"
    )

    log(
        "[config] cleanup pending: "
        + (
            "DISABILITATO"
            if args.no_cleanup_pending
            else "ATTIVO"
        )
    )

    log()
    disposable = set()
    source_stats = {}
    successful_sources = 0

    for url in SOURCES:
        log(
            f"[source] scarico {url}"
        )

        try:
            domains = download_source(
                url
            )

            source_stats[
                url
            ] = len(domains)

            successful_sources += 1

            disposable |= domains

            log(
                f"[source] trovati "
                f"{len(domains):,} domini"
            )

        except Exception as exc:
            source_stats[
                url
            ] = (
                "ERROR: "
                + str(exc)
            )

            log(
                f"[source] ERRORE: "
                f"{exc}"
            )

    if successful_sources == 0:
        raise RuntimeError(
            "Tutte le sorgenti sono fallite. "
            "Nessuna modifica effettuata."
        )

    disposable -= allowlist

    log()

    log(
        f"[source] unione finale: "
        f"{len(disposable):,} "
        f"domini disposable"
    )
    log()
    log(
        "[mastodon] leggo EmailDomainBlock..."
    )

    rows = get_mastodon_blocks(
        mastodon_dir
    )

    approvals = [
        row
        for row in rows
        if row[2] is True
    ]

    log(
        f"[mastodon] record totali: "
        f"{len(rows):,}"
    )

    log(
        f"[mastodon] approval: "
        f"{len(approvals):,}"
    )
    converted = 0

    if (
        approvals
        and not args.no_force_deny
    ):
        if args.dry_run:
            log(
                f"[DRY-RUN] convertirei "
                f"{len(approvals):,} "
                f"approval -> deny"
            )

        else:
            converted = (
                force_everything_to_deny(
                    mastodon_dir
                )
            )

            log(
                f"[mastodon] convertiti "
                f"approval -> deny: "
                f"{converted:,}"
            )

            rows = get_mastodon_blocks(
                mastodon_dir
            )
    existing = set()

    for row in rows:
        domain = normalize_domain(
            row[1]
        )

        if domain:
            existing.add(domain)

    missing_domains = sorted(
        disposable - existing
    )

    log()

    log(
        f"[domains] già presenti: "
        f"{len(existing):,}"
    )

    log(
        f"[domains] disposable "
        f"da aggiungere: "
        f"{len(missing_domains):,}"
    )

    batches = list(
        chunks(
            missing_domains,
            args.batch_size,
        )
    )

    domains_added = 0

    for number, batch in enumerate(
        batches,
        1,
    ):
        log(
            f"[domains] batch "
            f"{number}/"
            f"{len(batches)} "
            f"({len(batch):,})"
        )

        tootctl_add(
            mastodon_dir,
            batch,
            dry_run=args.dry_run,
        )

        domains_added += len(batch)
    cleanup = {
        "checked": 0,
        "matched": 0,
        "matched_email": 0,
        "matched_username": 0,
        "deleted": 0,
        "errors": [],
    }

    if not args.no_cleanup_pending:
        log()
        log(
            "[pending] controllo registrazioni "
            "in attesa..."
        )

        cleanup = cleanup_blocked_pending(
            mastodon_dir,
            dry_run=args.dry_run,
        )

        log(
            f"[pending] controllate : "
            f"{cleanup['checked']:,}"
        )

        log(
            f"[pending] bloccate    : "
            f"{cleanup['matched']:,}"
        )

        log(
            f"[pending] per email   : "
            f"{cleanup['matched_email']:,}"
        )

        log(
            f"[pending] per usr_    : "
            f"{cleanup['matched_username']:,}"
        )

        if args.dry_run:
            log(
                f"[pending] da eliminare: "
                f"{cleanup['matched']:,} "
                f"(DRY-RUN)"
            )

        else:
            log(
                f"[pending] eliminate  : "
                f"{cleanup['deleted']:,}"
            )

        if cleanup["errors"]:
            log(
                f"[pending] ERRORI     : "
                f"{len(cleanup['errors']):,}"
            )

            for error in (
                cleanup["errors"][:20]
            ):
                log(
                    "    "
                    f"user={error['id']} "
                    f"{error.get('username', '')} "
                    f"{error['error']}"
                )
    mx_to_domains = defaultdict(set)
    domains_with_mx = 0
    domains_without_mx = 0

    if args.mx_mode == "report":
        all_domains = sorted(
            disposable
        )

        log()
        log(
            f"[mx] analizzo gli MX di "
            f"{len(all_domains):,} domini"
        )

        log(
            "[mx] nessun MX verrà bloccato"
        )

        log(
            f"[mx] workers: "
            f"{args.mx_workers}"
        )

        with concurrent.futures.ThreadPoolExecutor(
            max_workers=args.mx_workers
        ) as executor:

            futures = {
                executor.submit(
                    resolve_mx,
                    domain,
                    args.mx_timeout,
                ): domain

                for domain in all_domains
            }

            completed = 0

            for future in (
                concurrent.futures
                .as_completed(futures)
            ):
                try:
                    _, mx_records = (
                        future.result()
                    )

                except Exception:
                    mx_records = []

                if mx_records:
                    domains_with_mx += 1
                else:
                    domains_without_mx += 1

                domain = futures[
                    future
                ]

                for mx in mx_records:
                    mx_to_domains[
                        mx
                    ].add(domain)

                completed += 1

                if (
                    completed % 5000 == 0
                    or completed
                    == len(all_domains)
                ):
                    log(
                        f"[mx] analizzati "
                        f"{completed:,}/"
                        f"{len(all_domains):,}"
                    )

    top_mx = sorted(
        [
            (
                mx,
                len(domains),
            )
            for (
                mx,
                domains,
            ) in mx_to_domains.items()
        ],
        key=lambda item: (
            -item[1],
            item[0],
        ),
    )[:100]

    if args.mx_mode == "report":
        log()

        log(
            f"[mx] domini con MX: "
            f"{domains_with_mx:,}"
        )

        log(
            f"[mx] senza MX / lookup fallito: "
            f"{domains_without_mx:,}"
        )

        log(
            f"[mx] MX unici osservati: "
            f"{len(mx_to_domains):,}"
        )

        log()
        log(
            "[mx] principali MX osservati "
            "(SOLO REPORT):"
        )

        for mx, count in top_mx:
            log(
                f"    {count:7,d} "
                f"domini -> {mx}"
            )
    if args.dry_run:
        final_rows = rows

    else:
        log()
        log(
            "[verify] rileggo "
            "EmailDomainBlock..."
        )

        final_rows = get_mastodon_blocks(
            mastodon_dir
        )

    final_deny = sum(
        1
        for row in final_rows
        if row[2] is False
    )

    final_approval = sum(
        1
        for row in final_rows
        if row[2] is True
    )
    report = {
        "version":
            VERSION,

        "timestamp":
            int(time.time()),

        "dry_run":
            args.dry_run,

        "sources":
            source_stats,

        "successful_sources":
            successful_sources,

        "disposable_domains":
            len(disposable),

        "allowlist":
            len(allowlist),

        "existing_before":
            len(rows),

        "approval_before":
            len(approvals),

        "approval_converted":
            converted,

        "missing_domains":
            len(missing_domains),

        "domains_added_or_planned":
            domains_added,

        "pending_checked":
            cleanup["checked"],

        "pending_blocked":
            cleanup["matched"],

        "pending_matched_email":
            cleanup["matched_email"],

        "pending_matched_username":
            cleanup["matched_username"],

        "pending_deleted":
            cleanup["deleted"],

        "pending_errors":
            cleanup["errors"],

        "mx_mode":
            args.mx_mode,

        "mx_autoblock":
            False,

        "domains_with_mx":
            domains_with_mx,

        "domains_without_mx":
            domains_without_mx,

        "mx_unique":
            len(mx_to_domains),

        "final_deny":
            final_deny,

        "final_approval":
            final_approval,

        "top_mx": [
            {
                "mx":
                    mx,

                "disposable_domains":
                    count,

                "action":
                    "REPORT_ONLY",
            }

            for mx, count in top_mx
        ],
    }

    save_json_atomic(
        args.report,
        report,
    )

    save_json_atomic(
        args.state,
        {
            "version":
                VERSION,

            "last_run":
                int(time.time()),

            "disposable_domains":
                len(disposable),

            "pending_deleted":
                cleanup["deleted"],

            "mx_unique":
                len(mx_to_domains),

            "mx_autoblock":
                False,

            "last_report":
                args.report,
        },
    )
    log()
    log(
        "============================================"
    )

    if args.dry_run:
        log(
            " RISULTATO DRY-RUN"
        )
    else:
        log(
            " RISULTATO"
        )

    log(
        "============================================"
    )

    log(
        f"Disposable aggregati : "
        f"{len(disposable):,}"
    )

    log(
        f"Domini nuovi         : "
        f"{domains_added:,}"
    )

    log(
        f"Pending controllati  : "
        f"{cleanup['checked']:,}"
    )

    log(
        f"Pending bloccati     : "
        f"{cleanup['matched']:,}"
    )

    log(
        f"  per email          : "
        f"{cleanup['matched_email']:,}"
    )

    log(
        f"  per usr_           : "
        f"{cleanup['matched_username']:,}"
    )

    log(
        f"Pending eliminati    : "
        f"{cleanup['deleted']:,}"
    )

    log(
        f"Errori pending       : "
        f"{len(cleanup['errors']):,}"
    )

    log(
        f"MX osservati         : "
        f"{len(mx_to_domains):,}"
    )

    log(
        "MX bloccati          : 0 "
        "(intenzionale)"
    )

    log(
        f"Deny finali          : "
        f"{final_deny:,}"
    )

    log(
        f"Approval finali      : "
        f"{final_approval:,}"
    )

    log(
        f"Report               : "
        f"{args.report}"
    )

    log()


if __name__ == "__main__":
    try:
        main()

    except KeyboardInterrupt:
        log()
        log("Interrotto.")
        sys.exit(130)

    except Exception as exc:
        log()
        log("ERRORE FATALE:")
        log(str(exc))
        raise