mosdns/tools/init.go

/*
 * Copyright (C) 2020-2022, IrineSistiana
 *
 * This file is part of mosdns.
 *
 * mosdns is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * mosdns is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

package tools

import (
	"fmt"
	"os"

	"github.com/IrineSistiana/mosdns/v5/coremain"
	"github.com/spf13/cobra"
)

func init() {
	probeCmd := &cobra.Command{
		Use:   "probe",
		Short: "Run some server tests.",
	}
	probeCmd.AddCommand(
		newConnReuseCmd(),
		newIdleTimeoutCmd(),
		newPipelineCmd(),
	)
	coremain.AddSubCmd(probeCmd)

	configCmd := &cobra.Command{
		Use:   "config",
		Short: "Tools that can generate/convert mosdns config file.",
	}
	configCmd.AddCommand(newGenCmd(), newConvCmd())
	coremain.AddSubCmd(configCmd)

	// 添加 init 命令用于快速初始化配置
	initCmd := &cobra.Command{
		Use:   "init",
		Short: "Initialize mosdns configuration and directories.",
		Long:  "Create default config.yaml and necessary directories for quick deployment on any server.",
		RunE:  runInit,
	}
	var forceFlag bool
	initCmd.Flags().BoolVarP(&forceFlag, "force", "f", false, "强制覆盖已存在的配置文件")
	coremain.AddSubCmd(initCmd)
}

// runInit 执行初始化操作
func runInit(cmd *cobra.Command, args []string) error {
	force, _ := cmd.Flags().GetBool("force")

	fmt.Println("========================================")
	fmt.Println("  🚀 MosDNS 初始化向导")
	fmt.Println("========================================")
	fmt.Println()

	// 1. 检查并创建配置文件
	configFile := "config.yaml"
	if err := createConfigFile(configFile, force); err != nil {
		return err
	}

	// 2. 创建必要的目录结构
	if err := createDirectories(); err != nil {
		return err
	}

	// 3. 创建示例数据文件（如果不存在）
	if err := createDataFiles(); err != nil {
		return err
	}

	// 4. 显示完成信息
	showCompletionInfo()

	return nil
}

// createConfigFile 创建默认配置文件
func createConfigFile(filename string, force bool) error {
	// 检查文件是否已存在
	if _, err := os.Stat(filename); err == nil && !force {
		fmt.Printf("⚠️  配置文件已存在: %s\n", filename)
		fmt.Println("   使用 --force 或 -f 参数强制覆盖")
		return nil
	}

	// 生成默认配置内容
	configContent := `# ========================================
# MosDNS 配置文件 - 智能 DNS 服务器
# 自动生成时间: $(date)
# ========================================

# 日志配置
log:
  level: info      # 日志级别: debug, info, warn, error
  file: ""         # 日志文件路径（空表示输出到控制台）

# API 管理接口配置
api:
  http: "0.0.0.0:8080"  # API 监听地址和端口

# Web 管理界面配置
web:
  http: "0.0.0.0:5555"  # Web UI 监听地址和端口

# 插件配置
plugins:
  # ========================================
  # 1. 数据源插件
  # ========================================
  
  # 中国 IP 地址库（用于智能分流）
  - tag: geoip_cn
    type: ip_set
    args:
      files:
        - "./data/chn_ip.txt"

  # 中国域名列表（用于智能分流）
  - tag: geosite_cn
    type: domain_set
    args:
      files:
        - "./data/geosite_china-list.txt"

  # ========================================
  # 2. DNS 上游服务器
  # ========================================
  
  # 国内 DNS 上游（用于解析国内域名）
  - tag: forward_local
    type: forward
    args:
      concurrent: 2  # 并发查询数量
      upstreams:
        - addr: "223.5.5.5"       # 阿里云 DNS
        - addr: "119.29.29.29"    # 腾讯云 DNS

  # 国外 DNS 上游（用于解析国外域名）
  - tag: forward_remote
    type: forward
    args:
      concurrent: 2
      upstreams:
        - addr: "https://1.1.1.1/dns-query"  # Cloudflare DoH
          enable_http3: false
        - addr: "https://8.8.8.8/dns-query"  # Google DoH
          enable_http3: false

  # ========================================
  # 3. 缓存插件
  # ========================================
  
  - tag: main_cache
    type: cache
    args:
      size: 100000           # 缓存条目数量
      lazy_cache_ttl: 86400  # 惰性缓存 TTL（秒）
      dump_file: "./cache.dump"      # 缓存持久化文件
      dump_interval: 3600    # 缓存保存间隔（秒）

  # ========================================
  # 4. 主执行序列
  # ========================================
  
  - tag: main
    type: sequence
    args:
      # 先查缓存
      - exec: $main_cache
      
      # 如果是国内域名，使用国内 DNS
      - matches:
          - qname $geosite_cn
        exec: $forward_local
      
      # 其他域名使用国外 DNS
      - exec: $forward_remote
      
      # 将结果存入缓存
      - matches:
          - has_resp
        exec: $main_cache

  # ========================================
  # 5. DNS 服务器
  # ========================================
  
  # UDP 服务器
  - tag: udp_server
    type: udp_server
    args:
      entry: main              # 入口执行序列
      listen: ":53"            # 监听端口（需要 root 权限）

  # TCP 服务器
  - tag: tcp_server
    type: tcp_server
    args:
      entry: main
      listen: ":53"
`

	// 写入文件
	if err := os.WriteFile(filename, []byte(configContent), 0644); err != nil {
		return fmt.Errorf("创建配置文件失败: %w", err)
	}

	fmt.Printf("✅ 配置文件已创建: %s\n", filename)
	return nil
}

// createDirectories 创建必要的目录结构
func createDirectories() error {
	dirs := []string{
		"./data",           // 数据文件目录
		"./config.d",       // 配置文件目录
		"./config.d/rules", // 规则文件目录
		"./logs",           // 日志目录（可选）
	}

	fmt.Println()
	fmt.Println("📁 创建目录结构...")

	for _, dir := range dirs {
		if err := os.MkdirAll(dir, 0755); err != nil {
			return fmt.Errorf("创建目录失败 %s: %w", dir, err)
		}
		fmt.Printf("   ✅ %s\n", dir)
	}

	return nil
}

// createDataFiles 创建示例数据文件
func createDataFiles() error {
	fmt.Println()
	fmt.Println("📄 检查数据文件...")

	dataFiles := map[string]string{
		"./data/chn_ip.txt": `# 中国 IP 地址段（示例）
# 请从以下地址下载完整列表：
# https://github.com/17mon/china_ip_list

# 示例 IP 段
1.0.1.0/24
1.0.2.0/23
`,
		"./data/geosite_china-list.txt": `# 中国常见域名列表（示例）
# 请从以下地址下载完整列表：
# https://github.com/felixonmars/dnsmasq-china-list

# 示例域名
domain:baidu.com
domain:qq.com
domain:taobao.com
domain:tmall.com
domain:jd.com
`,
	}

	for file, content := range dataFiles {
		// 如果文件已存在且不为空，跳过
		if stat, err := os.Stat(file); err == nil && stat.Size() > 0 {
			fmt.Printf("   ⏭️  已存在: %s\n", file)
			continue
		}

		// 创建示例文件
		if err := os.WriteFile(file, []byte(content), 0644); err != nil {
			return fmt.Errorf("创建数据文件失败 %s: %w", file, err)
		}
		fmt.Printf("   ✅ 已创建: %s\n", file)
	}

	return nil
}

// showCompletionInfo 显示完成信息和后续步骤
func showCompletionInfo() {
	wd, _ := os.Getwd()

	fmt.Println()
	fmt.Println("========================================")
	fmt.Println("  🎉 初始化完成！")
	fmt.Println("========================================")
	fmt.Println()
	fmt.Println("📂 工作目录:", wd)
	fmt.Println()
	fmt.Println("📋 已创建的文件和目录:")
	fmt.Println("   config.yaml           - 主配置文件")
	fmt.Println("   data/                 - 数据文件目录")
	fmt.Println("   config.d/rules/       - 规则文件目录")
	fmt.Println()
	fmt.Println("⚠️  重要提示:")
	fmt.Println("   1. 数据文件为示例文件，建议下载完整的 CN IP 和域名列表")
	fmt.Println("      - CN IP:   https://github.com/17mon/china_ip_list")
	fmt.Println("      - CN 域名: https://github.com/felixonmars/dnsmasq-china-list")
	fmt.Println()
	fmt.Println("   2. 默认端口 53 需要 root 权限，可以修改为其他端口（如 5310）")
	fmt.Println()
	fmt.Println("🚀 启动服务:")
	fmt.Println()
	fmt.Println("   # 开发模式（非 root）")
	fmt.Println("   sed -i 's/:53/:5310/g' config.yaml")
	fmt.Println("   ./mosdns-linux-amd64 start -c config.yaml")
	fmt.Println()
	fmt.Println("   # 生产模式（需要 root）")
	fmt.Println("   sudo ./mosdns-linux-amd64 start -c config.yaml")
	fmt.Println()
	fmt.Println("🌐 管理界面:")
	fmt.Println("   Web UI:  http://localhost:5555")
	fmt.Println("   API:     http://localhost:8080")
	fmt.Println()
	fmt.Println("========================================")
	fmt.Println()
}