ss/scientific_surfing/macos_service_wrapper.py

#!/usr/bin/env python3
"""
macOS service wrapper for mihomo.

This wrapper handles DNS configuration when starting/stopping the mihomo service.
"""

import sys
import subprocess
import signal
import os
from pathlib import Path
from typing import Optional, List, Dict


class MacOSDNSManager:
    """Manages macOS DNS settings."""

    @staticmethod
    def get_network_services() -> List[str]:
        """Get list of network services."""
        try:
            result = subprocess.run(
                ["networksetup", "-listallnetworkservices"],
                capture_output=True,
                text=True,
                check=True
            )
            # Skip first line (header)
            services = [line.strip() for line in result.stdout.split('\n')[1:] if line.strip() and not line.startswith('*')]
            return services
        except subprocess.CalledProcessError as e:
            print(f"Failed to list network services: {e}", file=sys.stderr)
            return []

    @staticmethod
    def get_current_dns(service: str) -> List[str]:
        """
        Get current DNS servers for a network service.

        Args:
            service: Network service name

        Returns:
            List[str]: List of DNS server addresses
        """
        try:
            result = subprocess.run(
                ["networksetup", "-getdnsservers", service],
                capture_output=True,
                text=True,
                check=True
            )
            if result.returncode:
                return ["Empty"]
            dns_servers = [line.strip() for line in result.stdout.split('\n') if line.strip() and " " not in line.strip()]
            return dns_servers
        except subprocess.CalledProcessError:
            return ["Empty"]

    @staticmethod
    def set_dns(service: str, dns_servers: List[str]) -> bool:
        """
        Set DNS servers for a network service.

        Args:
            service: Network service name
            dns_servers: List of DNS server addresses (empty list to reset to DHCP)

        Returns:
            bool: True if successful, False otherwise
        """
        try:
            if not dns_servers:
                # Reset to DHCP/automatic
                subprocess.run(
                    ["networksetup", "-setdnsservers", service, "empty"],
                    capture_output=True,
                    check=True
                )
            else:
                subprocess.run(
                    ["networksetup", "-setdnsservers", service] + dns_servers,
                    capture_output=True,
                    check=True
                )
            return True
        except subprocess.CalledProcessError as e:
            print(f"Failed to set DNS for {service}: {e}", file=sys.stderr)
            return False


class MihomoServiceWrapper:
    """Wrapper for mihomo service on macOS."""

    def __init__(self, executable: str, args: List[str], config_dir: Path):
        """
        Initialize the service wrapper.

        Args:
            executable: Path to mihomo executable
            args: Command line arguments
            config_dir: Configuration directory
        """
        self.executable = executable
        self.args = args
        self.config_dir = config_dir
        self.process: Optional[subprocess.Popen] = None
        self.dns_manager = MacOSDNSManager()
        self.original_dns: Dict[str, List[str]] = {}
        self.shutting_down = False

    def start(self) -> bool:
        """
        Start the mihomo service and configure DNS.

        Returns:
            bool: True if successful, False otherwise
        """
        try:
            # Start mihomo process
            cmd = [self.executable] + self.args
            self.process = subprocess.Popen(
                cmd,
                stdout=sys.stdout,
                stderr=sys.stderr,
                cwd=str(self.config_dir),
            )

            print(f"Started mihomo service (PID: {self.process.pid})", file=sys.stderr)

            # Configure DNS for network services
            dns_server = "198.18.0.2"
            services = self.dns_manager.get_network_services()

            for service in services:
                # Save original DNS settings
                original_dns = self.dns_manager.get_current_dns(service)
                self.original_dns[service] = original_dns

                # Set new DNS
                if self.dns_manager.set_dns(service, [dns_server]):
                    print(f"Set DNS for {service}: {dns_server}", file=sys.stderr)
                else:
                    print(f"Warning: Failed to set DNS for {service}", file=sys.stderr)

            return True
        except Exception as e:
            print(f"Failed to start service: {e}", file=sys.stderr)
            return False

    def stop(self) -> bool:
        """
        Stop the mihomo service and restore DNS settings.

        Returns:
            bool: True if successful, False otherwise
        """
        try:
            # Restore original DNS settings
            for service, dns_servers in self.original_dns.items():
                if self.dns_manager.set_dns(service, dns_servers):
                    if dns_servers:
                        print(f"Restored DNS for {service}: {', '.join(dns_servers)}", file=sys.stderr)
                    else:
                        print(f"Reset DNS for {service} to automatic", file=sys.stderr)
                else:
                    print(f"Warning: Failed to restore DNS for {service}", file=sys.stderr)

            # Stop mihomo process
            if self.process:
                self.process.terminate()
                try:
                    self.process.wait(timeout=5)
                except subprocess.TimeoutExpired:
                    self.process.kill()
                    self.process.wait()
                print(f"Stopped mihomo service (PID: {self.process.pid})", file=sys.stderr)

            return True
        except Exception as e:
            print(f"Failed to stop service: {e}", file=sys.stderr)
            return False

    def run(self) -> None:
        """Run the service wrapper (blocking)."""
        self.shutting_down = False

        # Setup signal handlers
        signal.signal(signal.SIGTERM, self._handle_signal)
        signal.signal(signal.SIGINT, self._handle_signal)

        # Start service
        if not self.start():
            sys.exit(1)

        print("Service started, press Ctrl+C to stop", file=sys.stderr)

        # Wait for process to exit
        try:
            if self.process:
                # Use poll with timeout to allow signal handling
                while self.process.poll() is None and not self.shutting_down:
                    try:
                        self.process.wait(timeout=0.5)
                    except subprocess.TimeoutExpired:
                        continue

                if self.process.poll() is not None:
                    print(f"Process exited with code: {self.process.returncode}", file=sys.stderr)
        except KeyboardInterrupt:
            print("\nReceived KeyboardInterrupt", file=sys.stderr)
        finally:
            if not self.shutting_down:
                self.stop()

    def _handle_signal(self, signum: int, _frame) -> None:
        """Handle termination signals."""
        if self.shutting_down:
            # Force kill on second signal
            print(f"Force killing...", file=sys.stderr)
            if self.process:
                self.process.kill()
            sys.exit(1)

        print(f"Received signal {signum}, shutting down...", file=sys.stderr)
        self.shutting_down = True
        self.stop()
        sys.exit(0)


def main() -> None:
    """Main entry point."""
    if len(sys.argv) < 3:
        print("Usage: macos_service_wrapper.py <executable> <config_dir> [args...]", file=sys.stderr)
        sys.exit(1)

    executable = sys.argv[1]
    config_dir = Path(sys.argv[2])
    args = sys.argv[3:] if len(sys.argv) > 3 else []

    # Validate inputs
    if not Path(executable).exists():
        print(f"Error: Executable not found: {executable}", file=sys.stderr)
        sys.exit(1)

    if not config_dir.exists():
        print(f"Error: Config directory not found: {config_dir}", file=sys.stderr)
        sys.exit(1)

    # Create and run wrapper
    wrapper = MihomoServiceWrapper(executable, args, config_dir)
    wrapper.run()


if __name__ == "__main__":
    main()