main.py

"""
Клиент DNS Lookup API Atlorium — DNS-записи домена и аудит защиты от спуфинга.

Запуск (работает сразу, без регистрации — на демо-ключе):
    pip install -r requirements.txt
    python main.py
    python main.py example.com

Боевой ключ: получить на https://atlorium.com и положить в переменную окружения
ATLORIUM_API_KEY. Код при этом не меняется.
"""

import os
import sys
import time
from dataclasses import dataclass, field

import requests

# Публичный демо-ключ. С ним API отвечает правдоподобными МОКАМИ (не реальными
# записями DNS) — чтобы можно было встроить и протестировать интеграцию до оплаты.
# Ответы детерминированы: один и тот же домен всегда даёт один и тот же результат,
# поэтому на них можно писать стабильные тесты.
SANDBOX_KEY = "ak_sandbox_demo_mockdata_v1"

API_KEY = os.environ.get("ATLORIUM_API_KEY", SANDBOX_KEY)
BASE_URL = os.environ.get("ATLORIUM_BASE_URL", "https://atlorium.com")

TIMEOUT = 30

# 429 — не поломка, а штатный лимит. Ждём и повторяем.
RETRY_DELAY = 20
MAX_RETRIES = 2

# Потолок ожидания. Исчерпав ЧАСОВОЙ лимит, сервер честно просит подождать
# 40+ минут — и клиент, слепо доверяющий заголовку Retry-After, зависнет на все
# 40 минут (а в CI просто съест бюджет джоба). Дольше потолка не ждём: сообщаем,
# что квота исчерпана, и выходим.
MAX_RETRY_DELAY = 120


class AtloriumError(RuntimeError):
    """Ошибка API. Код HTTP разложен в человекочитаемую причину."""

    REASONS = {
        400: "Домен не передан или задан неверно (нужно чистое имя, без https:// и без IP)",
        401: "API-ключ отсутствует, просрочен или недействителен",
        402: "Недостаточно кредитов на балансе — пополните на https://atlorium.com",
        429: "Превышен лимит запросов — повторите позже",
        500: "DNS-резолвер не ответил или превышен таймаут запроса",
        503: "Сервис временно недоступен "
             "(за сбой на своей стороне мы не списываем деньги)",
    }

    def __init__(self, status: int, body: str):
        reason = self.REASONS.get(status, "Неизвестная ошибка")
        super().__init__(f"HTTP {status}: {reason}. Ответ сервера: {body[:200]}")
        self.status = status


def retry_after(response: requests.Response) -> int:
    """Сколько ждать после 429. Ноль/мусор и слишком большие значения не берём на веру.

    Значение 0 (или мусор) означало бы «повторяй немедленно» — клиент ушёл бы в
    busy-loop и выжег остаток лимита за секунду. Значение в 40+ минут (так сервер
    отвечает на исчерпанный часовой лимит) означало бы «спи почти час» — этого мы
    тоже не делаем. Возвращаем 0, если ждать бессмысленно долго: вызывающий сдастся.
    """
    raw = response.headers.get("Retry-After", "")
    try:
        seconds = int(raw)
    except ValueError:
        seconds = 0

    if seconds <= 0:
        return RETRY_DELAY
    return seconds if seconds <= MAX_RETRY_DELAY else 0


def _get(path: str, params: dict) -> dict:
    for attempt in range(MAX_RETRIES + 1):
        response = requests.get(
            f"{BASE_URL}{path}",
            params={key: value for key, value in params.items() if value},
            headers={
                "Authorization": f"Bearer {API_KEY}",
                "Accept": "application/json",
            },
            timeout=TIMEOUT,
        )

        if response.status_code == 429 and attempt < MAX_RETRIES:
            delay = retry_after(response)
            if delay == 0:
                raise AtloriumError(429, "лимит запросов по IP исчерпан, повторите позже")
            print(f"  ... лимит запросов, пауза {delay} с", file=sys.stderr)
            time.sleep(delay)
            continue

        if not response.ok:
            raise AtloriumError(response.status_code, response.text)
        return response.json()

    raise AtloriumError(429, "лимит запросов не отпустил после повторов")


def lookup(domain: str, resolver: str | None = None) -> dict:
    """DNS-записи домена: A, AAAA, MX, TXT, CNAME, NS — одним запросом.

    resolver — IP-адрес резолвера (по умолчанию Cloudflare 1.1.1.1).
    """
    return _get("/api/dns/lookup", {"domain": domain, "resolver": resolver})


def check_propagation(domain: str, resolvers: str | None = None) -> dict:
    """Распространение записей по резолверам: одинаково ли они видят домен.

    Нужно после смены DNS: пока резолверы отдают разное, часть пользователей ещё
    ходит на старый адрес.

    resolvers — до 8 публичных IP через запятую; повторы схлопываются. Битый или
    внутренний адрес, как и список длиннее восьми, дают 400, а не молчаливое усечение.

    Тарифицируется ПОРЕЗОЛВЕРНО: N резолверов = N единиц (по умолчанию их 4).
    Поэтому в примере функция не вызывается — чтобы прогон не расходовал лишнего.
    """
    return _get("/api/dns/propagation", {"domain": domain, "resolvers": resolvers})


def records_of(response: dict, record_type: str) -> list[dict]:
    """Записи одного типа. Ответ приходит группами: [{type, records: [...]}, ...]."""
    for group in response.get("records") or []:
        if group.get("type") == record_type:
            return group.get("records") or []
    return []


# ── Применение данных: аудит защиты домена от подделки писем ──────────────────
# Список DNS-записей сам по себе — просто JSON. Ценность появляется, когда из него
# делают вывод. Ниже — проверка, может ли посторонний слать письма ОТ ИМЕНИ домена.
#
# Три кирпича защиты:
#   MX    — принимает ли домен входящую почту вообще;
#   SPF   — TXT-запись «v=spf1 ...»: кому разрешено слать от имени домена;
#   DMARC — TXT-запись на ПОДДОМЕНЕ _dmarc.<домен>: что делать с письмом, которое
#           проверку не прошло. Это отдельный второй запрос к API.
#
# Ровно два запроса за прогон: домен и _dmarc.домен.

STRICT = "-all"    # hard fail: чужому письму отказать
SOFT = "~all"      # softfail: пометить как подозрительное, но доставить
NEUTRAL = "?all"   # neutral: политики фактически нет
PERMISSIVE = "+all"  # разрешено всем — дыра


@dataclass
class Issue:
    level: str  # "!" — критично, "~" — слабое место, "i" — к сведению
    text: str


@dataclass
class MailSecurity:
    domain: str
    # Полный ответ по домену — чтобы напечатать записи, не делая третий запрос.
    root: dict = field(default_factory=dict)
    mx: list[str] = field(default_factory=list)
    spf: str | None = None
    spf_policy: str | None = None      # -all / ~all / ?all / +all / None
    dmarc: str | None = None
    dmarc_policy: str | None = None    # none / quarantine / reject / None
    verdict: str = ""
    issues: list[Issue] = field(default_factory=list)


def parse_spf_policy(spf: str) -> str | None:
    """Механизм «all» в конце SPF — это и есть политика. Без него по умолчанию neutral."""
    for policy in (STRICT, SOFT, NEUTRAL, PERMISSIVE):
        if policy in spf:
            return policy
    if " all" in f" {spf}":  # голое «all» без квалификатора == +all
        return PERMISSIVE
    return None


def parse_dmarc_policy(dmarc: str) -> str | None:
    """Тег p= в DMARC: none (только мониторинг) / quarantine (в спам) / reject (отбросить)."""
    for tag in dmarc.split(";"):
        key, _, value = tag.strip().partition("=")
        if key.strip().lower() == "p":
            return value.strip().lower()
    return None


def audit_mail_security(domain: str) -> MailSecurity:
    """Аудит защиты домена от спуфинга. Ровно 2 запроса к API."""
    result = MailSecurity(domain=domain)

    # Запрос 1 — записи самого домена: MX и TXT (в TXT живёт SPF).
    root = lookup(domain)
    result.root = root
    result.mx = [
        f"{record['value']} (приоритет {record.get('preference')})"
        if record.get("preference") is not None else record["value"]
        for record in records_of(root, "MX")
    ]

    txt_values = [record["value"] for record in records_of(root, "TXT")]
    result.spf = next((txt for txt in txt_values if txt.lower().startswith("v=spf1")), None)

    # Запрос 2 — DMARC лежит НЕ на домене, а на поддомене _dmarc.<домен>.
    # Это отдельная зона, поэтому и запрос отдельный.
    dmarc_zone = lookup(f"_dmarc.{domain}")
    dmarc_values = [record["value"] for record in records_of(dmarc_zone, "TXT")]
    result.dmarc = next((txt for txt in dmarc_values if txt.lower().startswith("v=dmarc1")), None)

    # ── MX: принимает ли домен почту ──────────────────────────────────────────
    if not result.mx:
        result.issues.append(Issue("i",
            "MX-записей нет — домен не принимает входящую почту. От спуфинга это не "
            "защищает: слать письма ОТ его имени всё равно можно, и SPF/DMARC нужны."))

    # ── SPF: кому разрешено слать от имени домена ─────────────────────────────
    if result.spf is None:
        spf_score = 0
        result.issues.append(Issue("!",
            "SPF не найден: письма от имени домена может слать кто угодно — "
            "принимающий сервер не с чем сверить отправителя"))
    else:
        result.spf_policy = parse_spf_policy(result.spf)
        if result.spf_policy == STRICT:
            spf_score = 3
        elif result.spf_policy == SOFT:
            spf_score = 2
            result.issues.append(Issue("~",
                "SPF ~all (softfail): подделку помечают, но письмо всё равно доставляют. "
                "Строгая политика — -all"))
        elif result.spf_policy in (NEUTRAL, PERMISSIVE):
            spf_score = 1
            result.issues.append(Issue("!",
                f"SPF-политика {result.spf_policy}: запись есть, но разрешает всех — "
                "политики фактически нет"))
        else:
            spf_score = 1
            result.issues.append(Issue("!",
                "В SPF нет механизма all: по умолчанию действует neutral — "
                "политика не работает"))

    # ── DMARC: что делать с письмом, не прошедшим проверку ────────────────────
    if result.dmarc is None:
        dmarc_score = 0
        result.issues.append(Issue("!",
            f"DMARC не найден (_dmarc.{domain}): почтовый сервер получателя не знает, "
            "что делать с письмом, не прошедшим SPF/DKIM, — и обычно его доставляет"))
    else:
        result.dmarc_policy = parse_dmarc_policy(result.dmarc)
        if result.dmarc_policy == "reject":
            dmarc_score = 3
        elif result.dmarc_policy == "quarantine":
            dmarc_score = 2
            result.issues.append(Issue("~",
                "DMARC p=quarantine: подделки уходят в спам, но не отбрасываются. "
                "Максимальная защита — p=reject"))
        elif result.dmarc_policy == "none":
            dmarc_score = 1
            result.issues.append(Issue("~",
                "DMARC p=none: нарушения только логируются, письма не блокируются — "
                "это режим наблюдения, а не защиты"))
        else:
            dmarc_score = 0
            result.issues.append(Issue("!",
                "DMARC-запись есть, но тег p= отсутствует или неизвестен — "
                "запись недействительна"))

    # Итог по слабейшему звену: дыра в любом из двух механизмов делает домен
    # пригодным для подделки писем.
    score = min(spf_score, dmarc_score)
    if score >= 3:
        result.verdict = "ЗАЩИЩЁН — подделать отправителя не выйдет"
    elif score >= 1:
        result.verdict = "ЧАСТИЧНО ЗАЩИЩЁН — защита есть, но её можно обойти"
    else:
        result.verdict = "УЯЗВИМ — от имени домена можно слать поддельные письма"

    return result


def main() -> int:
    if API_KEY == SANDBOX_KEY:
        print("Демо-ключ: ответы сгенерированы (моки), не реальные записи DNS.\n")

    domain = sys.argv[1] if len(sys.argv) > 1 else "atlorium.com"

    try:
        # Внутри ровно два запроса к API: сам домен и _dmarc.<домен>.
        # Ответ по домену возвращается в audit.root — печатаем его без третьего запроса.
        audit = audit_mail_security(domain)
    except AtloriumError as error:
        print(f"Ошибка: {error}", file=sys.stderr)
        return 1

    root = audit.root
    print(f"DNS-записи {root.get('domain')} (резолвер {root.get('resolver')}):")
    for group in root.get("records") or []:
        for record in group.get("records") or []:
            preference = record.get("preference")
            prefix = f"{preference} " if preference is not None else ""
            print(f"  {group['type']:<6} {prefix}{record['value']}  (TTL {record.get('ttl')})")
    for error in root.get("errors") or []:
        print(f"  [!] {error}")

    print(f"\nАУДИТ ЗАЩИТЫ ОТ СПУФИНГА — {domain}")
    print(f"  MX:    {', '.join(audit.mx) if audit.mx else 'нет'}")
    print(f"  SPF:   {audit.spf or 'не найден'}")
    print(f"  DMARC: {audit.dmarc or f'не найден (_dmarc.{domain})'}")

    print(f"\nВЕРДИКТ: {audit.verdict}")
    if audit.issues:
        print("\nПРОБЛЕМЫ:")
        for issue in audit.issues:
            print(f"  [{issue.level}] {issue.text}")

    return 0


if __name__ == "__main__":
    raise SystemExit(main())