[DNS Rebinding] 가내수공업 네임 서버 구성기

from __future__ import annotations

import os
import signal
import threading
import time
from dataclasses import dataclass
from ipaddress import IPv4Address, ip_address
from typing import Any

from dnslib import A, CLASS, QTYPE, RCODE, RR
from dnslib.server import BaseResolver, DNSServer

def env(name: str, default: str) -> str:
    value = os.getenv(name)
    return value if value else default

def env_int(name: str, default: int) -> int:
    raw = env(name, str(default))

    try:
        return int(raw)
    except ValueError as exc:
        raise ValueError(
            f"{name} must be an integer: {raw!r}"
        ) from exc

def env_float(name: str, default: float) -> float:
    raw = env(name, str(default))

    try:
        return float(raw)
    except ValueError as exc:
        raise ValueError(
            f"{name} must be a number: {raw!r}"
        ) from exc

def normalize_dns_name(
    name: str,
    field_name: str,
) -> str:
    """
    DNS 이름을 소문자 FQDN 형식으로 정규화합니다.

    example.com -> example.com.
    """
    normalized = name.strip().lower()

    if not normalized:
        raise ValueError(
            f"{field_name} must not be empty"
        )

    if normalized.endswith("."):
        return normalized

    return f"{normalized}."

@dataclass
class Session:
    count: int
    last_seen: float

@dataclass(frozen=True)
class ClientInfo:
    source_ip: str
    source_port: int
    transport: str

@dataclass(frozen=True)
class AnswerSpec:
    value: str

    @classmethod
    def from_a(cls, value: str) -> AnswerSpec:
        return cls(value=str(IPv4Address(value)))

    @property
    def type_name(self) -> str:
        return "A"

    def create_rr(
        self,
        qname,
        ttl: int,
    ) -> RR:
        return RR(
            rname=qname,
            rtype=QTYPE.A,
            rclass=CLASS.IN,
            ttl=ttl,
            rdata=A(self.value),
        )

def identify_client(handler: Any) -> ClientInfo:
    """
    dnslib 요청 핸들러에서 클라이언트 주소를 추출합니다.
    """
    client_address = getattr(
        handler,
        "client_address",
        None,
    )

    if not client_address or len(client_address) < 2:
        raise ValueError(
            "Unable to determine client address"
        )

    source_ip = str(ip_address(client_address[0]))
    source_port = int(client_address[1])

    transport = str(
        getattr(handler, "protocol", "unknown")
    ).lower()

    return ClientInfo(
        source_ip=source_ip,
        source_port=source_port,
        transport=transport,
    )

def load_second_answer() -> AnswerSpec:
    """두 번째 A 응답에 사용할 IPv4 주소를 읽습니다."""
    return AnswerSpec.from_a(
        env("SECOND_IP", "127.0.0.1")
    )

class ResponseSwitcher:
    def __init__(
        self,
        first_answer: AnswerSpec,
        second_answer: AnswerSpec,
        reset_after: float,
    ) -> None:
        self.first_answer = first_answer
        self.second_answer = second_answer
        self.reset_after = reset_after

        # 서버 전체에서 하나의 질의 순서만 관리합니다.
        self._session: Session | None = None
        self._lock = threading.Lock()

    def next_answer(self) -> tuple[AnswerSpec, int]:
        """
        서버 전체 기준 첫 번째 A 질의에는 first_answer,
        두 번째 이후 A 질의에는 second_answer를 반환합니다.

        reset_after 동안 A 질의가 없으면 다시 첫 번째부터 시작합니다.
        """
        now = time.monotonic()

        with self._lock:
            if (
                self._session is None
                or now - self._session.last_seen > self.reset_after
            ):
                self._session = Session(
                    count=0,
                    last_seen=now,
                )

            query_number = self._session.count + 1

            if self._session.count == 0:
                answer = self.first_answer
            else:
                answer = self.second_answer

            self._session.count += 1
            self._session.last_seen = now

            return answer, query_number

class ZoneResolver(BaseResolver):
    def __init__(
        self,
        zone_name: str,
        switcher: ResponseSwitcher,
        ttl: int,
    ) -> None:
        self.zone_name = normalize_dns_name(
            zone_name,
            "ZONE_NAME",
        )
        self.switcher = switcher
        self.ttl = ttl

    def _is_in_zone(self, qname: str) -> bool:
        """
        설정된 존 자체와 하위 도메인만 허용합니다.

        ZONE_NAME=example.com인 경우:

        example.com.       허용
        www.example.com.   허용
        api.example.com.   허용
        google.com.        거절
        """
        return (
            qname == self.zone_name
            or qname.endswith(f".{self.zone_name}")
        )

    def resolve(self, request, handler):
        reply = request.reply()
        reply.header.ra = 0

        if len(request.questions) != 1:
            reply.header.rcode = RCODE.FORMERR
            return reply

        question = request.q
        qname = str(question.qname).lower()

        if not self._is_in_zone(qname):
            reply.header.rcode = RCODE.REFUSED
            return reply

        reply.header.aa = 1

        # A/IN 질의에 대해서만 전환 응답을 반환합니다.
        if (
            question.qtype != QTYPE.A
            or question.qclass != CLASS.IN
        ):
            return reply

        try:
            client = identify_client(handler)
        except (ValueError, TypeError) as exc:
            print(
                f"client identification failed: {exc}",
                flush=True,
            )
            reply.header.rcode = RCODE.SERVFAIL
            return reply

        answer, query_number = self.switcher.next_answer()

        reply.add_answer(
            answer.create_rr(
                qname=question.qname,
                ttl=self.ttl,
            )
        )

        print(
            f"client_ip={client.source_ip} "
            f"client_port={client.source_port} "
            f"transport={client.transport} "
            f"name={qname} "
            f"query_number={query_number} "
            f"answer_type={answer.type_name} "
            f"answer={answer.value} "
            f"ttl={self.ttl}",
            flush=True,
        )

        return reply

def main() -> None:
    dns_host = env(
        "DNS_HOST",
        "127.0.0.1",
    )

    dns_port = env_int(
        "DNS_PORT",
        53,
    )

    zone_name = env(
        "ZONE_NAME",
        "switch.lab.test",
    )

    first_ip = env(
        "FIRST_IP",
        "192.0.2.10",
    )

    second_answer = load_second_answer()

    reset_seconds = env_float(
        "RESET_SECONDS",
        30.0,
    )

    # 즉시 재질의가 발생하도록 기본 TTL은 0입니다.
    dns_ttl = env_int(
        "DNS_TTL",
        0,
    )

    if not 1 <= dns_port <= 65535:
        raise ValueError(
            "DNS_PORT must be between 1 and 65535"
        )

    if reset_seconds <= 0:
        raise ValueError(
            "RESET_SECONDS must be positive"
        )

    if dns_ttl < 0:
        raise ValueError(
            "DNS_TTL must be zero or greater"
        )

    first_answer = AnswerSpec.from_a(first_ip)

    switcher = ResponseSwitcher(
        first_answer=first_answer,
        second_answer=second_answer,
        reset_after=reset_seconds,
    )

    resolver = ZoneResolver(
        zone_name=zone_name,
        switcher=switcher,
        ttl=dns_ttl,
    )

    udp_server = DNSServer(
        resolver,
        port=dns_port,
        address=dns_host,
        tcp=False,
    )

    tcp_server = DNSServer(
        resolver,
        port=dns_port,
        address=dns_host,
        tcp=True,
    )

    stop_event = threading.Event()

    def request_shutdown(signum, frame) -> None:
        del signum, frame
        stop_event.set()

    signal.signal(
        signal.SIGINT,
        request_shutdown,
    )

    signal.signal(
        signal.SIGTERM,
        request_shutdown,
    )

    udp_server.start_thread()
    tcp_server.start_thread()

    print(
        f"DNS UDP/TCP listening on "
        f"{dns_host}:{dns_port}; "
        f"zone={resolver.zone_name}; "
        f"first={first_answer.type_name}:"
        f"{first_answer.value}; "
        f"second={second_answer.type_name}:"
        f"{second_answer.value}; "
        f"ttl={dns_ttl}; "
        f"reset={reset_seconds}s",
        flush=True,
    )

    try:
        while not stop_event.wait(timeout=1):
            pass
    finally:
        udp_server.stop()
        tcp_server.stop()

if __name__ == "__main__":
    main()