fscan/Common/Flag.go

package common

import (
	"flag"
	"fmt"
	"os"
	"strings"

	"github.com/fatih/color"
	"github.com/shadow1ng/fscan/common/config"
	"github.com/shadow1ng/fscan/common/i18n"
)

// Flag专用变量 (只在Flag.go中使用的变量直接定义在这里)
var (
	ExcludeHosts string
	Ports        string
	ExcludePorts string
	AddPorts     string
	HostsFile    string
	PortsFile    string
	
	ModuleThreadNum   int
	GlobalTimeout     int64
	EnableFingerprint bool
	
	AddUsers      string
	AddPasswords  string
	UsersFile     string
	PasswordsFile string
	HashFile      string
	HashValue     string
	Domain        string
	SshKeyPath    string
	
	TargetURL    string
	URLsFile     string
	Cookie       string
	WebTimeout   int64
	UserAgent    string
	Accept       string
	
	PocPath        string
	PocFull        bool
	DnsLog         bool
	PocNum         int
	DisablePocScan bool
	
	RedisFile         string
	RedisShell        string
	DisableRedis      bool
	RedisWritePath    string
	RedisWriteContent string
	RedisWriteFile    string
	
	DisableBrute bool
	MaxRetries   int
	
	DisableSave   bool
	Silent        bool
	DisableProgress bool
	
	Shellcode string
	
	// Parse.go 使用的变量
	HostPort   []string
	URLs       []string
	HashValues []string
	HashBytes  [][]byte
)

// Pocinfo POC信息变量
var Pocinfo config.PocInfo

func Banner() {
	// 定义暗绿色系
	colors := []color.Attribute{
		color.FgGreen,   // 基础绿
		color.FgHiGreen, // 亮绿
	}

	lines := []string{
		"   ___                              _    ",
		"  / _ \\     ___  ___ _ __ __ _  ___| | __ ",
		" / /_\\/____/ __|/ __| '__/ _` |/ __| |/ /",
		"/ /_\\\\_____\\__ \\ (__| | | (_| | (__|   <    ",
		"\\____/     |___/\\___|_|  \\__,_|\\___|_|\\_\\   ",
	}

	// 获取最长行的长度
	maxLength := 0
	for _, line := range lines {
		if len(line) > maxLength {
			maxLength = len(line)
		}
	}

	// 创建边框
	topBorder := "┌" + strings.Repeat("─", maxLength+2) + "┐"
	bottomBorder := "└" + strings.Repeat("─", maxLength+2) + "┘"

	// 打印banner
	fmt.Println(topBorder)

	for lineNum, line := range lines {
		fmt.Print("│ ")
		// 使用对应的颜色打印每个字符
		c := color.New(colors[lineNum%2])
		c.Print(line)
		// 补齐空格
		padding := maxLength - len(line)
		fmt.Printf("%s │\n", strings.Repeat(" ", padding))
	}

	fmt.Println(bottomBorder)

	// 打印版本信息
	c := color.New(colors[1])
	c.Printf("      Fscan Version: %s\n\n", version)
}

// Flag 解析命令行参数并配置扫描选项
func Flag(Info *HostInfo) {
	Banner()

	// 预处理语言设置 - 在定义flag之前检查lang参数
	preProcessLanguage()

	// ═════════════════════════════════════════════════
	// 目标配置参数
	// ═════════════════════════════════════════════════
	flag.StringVar(&Info.Host, "h", "", i18n.GetText("flag_host"))
	flag.StringVar(&ExcludeHosts, "eh", "", i18n.GetText("flag_exclude_hosts"))
	flag.StringVar(&Ports, "p", MainPorts, i18n.GetText("flag_ports"))
	flag.StringVar(&ExcludePorts, "ep", "", i18n.GetText("flag_exclude_ports"))
	flag.StringVar(&HostsFile, "hf", "", i18n.GetText("flag_hosts_file"))
	flag.StringVar(&PortsFile, "pf", "", i18n.GetText("flag_ports_file"))

	// ═════════════════════════════════════════════════
	// 扫描控制参数
	// ═════════════════════════════════════════════════
	flag.StringVar(&ScanMode, "m", "all", i18n.GetText("flag_scan_mode"))
	flag.IntVar(&ThreadNum, "t", 600, i18n.GetText("flag_thread_num"))
	flag.Int64Var(&Timeout, "time", 3, i18n.GetText("flag_timeout"))
	flag.IntVar(&ModuleThreadNum, "mt", 10, i18n.GetText("flag_module_thread_num"))
	flag.Int64Var(&GlobalTimeout, "gt", 180, i18n.GetText("flag_global_timeout"))
	// LiveTop 参数已移除，改为智能控制
	flag.BoolVar(&DisablePing, "np", false, i18n.GetText("flag_disable_ping"))
	flag.BoolVar(&EnableFingerprint, "fingerprint", false, i18n.GetText("flag_enable_fingerprint"))
	flag.BoolVar(&LocalMode, "local", false, i18n.GetText("flag_local_mode"))
	flag.BoolVar(&AliveOnly, "ao", false, i18n.GetText("flag_alive_only"))

	// ═════════════════════════════════════════════════
	// 认证与凭据参数
	// ═════════════════════════════════════════════════
	flag.StringVar(&Username, "user", "", i18n.GetText("flag_username"))
	flag.StringVar(&Password, "pwd", "", i18n.GetText("flag_password"))
	flag.StringVar(&AddUsers, "usera", "", i18n.GetText("flag_add_users"))
	flag.StringVar(&AddPasswords, "pwda", "", i18n.GetText("flag_add_passwords"))
	flag.StringVar(&UsersFile, "userf", "", i18n.GetText("flag_users_file"))
	flag.StringVar(&PasswordsFile, "pwdf", "", i18n.GetText("flag_passwords_file"))
	flag.StringVar(&HashFile, "hashf", "", i18n.GetText("flag_hash_file"))
	flag.StringVar(&HashValue, "hash", "", i18n.GetText("flag_hash_value"))
	flag.StringVar(&Domain, "domain", "", i18n.GetText("flag_domain"))      // SMB扫描用
	flag.StringVar(&SshKeyPath, "sshkey", "", i18n.GetText("flag_ssh_key")) // SSH扫描用

	// ═════════════════════════════════════════════════
	// Web扫描参数
	// ═════════════════════════════════════════════════
	flag.StringVar(&TargetURL, "u", "", i18n.GetText("flag_target_url"))
	flag.StringVar(&URLsFile, "uf", "", i18n.GetText("flag_urls_file"))
	flag.StringVar(&Cookie, "cookie", "", i18n.GetText("flag_cookie"))
	flag.Int64Var(&WebTimeout, "wt", 5, i18n.GetText("flag_web_timeout"))
	flag.StringVar(&HttpProxy, "proxy", "", i18n.GetText("flag_http_proxy"))
	flag.StringVar(&Socks5Proxy, "socks5", "", i18n.GetText("flag_socks5_proxy"))

	// ═════════════════════════════════════════════════
	// POC测试参数
	// ═════════════════════════════════════════════════
	flag.StringVar(&PocPath, "pocpath", "", i18n.GetText("flag_poc_path"))
	flag.StringVar(&Pocinfo.PocName, "pocname", "", i18n.GetText("flag_poc_name"))
	flag.BoolVar(&PocFull, "full", false, i18n.GetText("flag_poc_full"))
	flag.BoolVar(&DnsLog, "dns", false, i18n.GetText("flag_dns_log"))
	flag.IntVar(&PocNum, "num", 20, i18n.GetText("flag_poc_num"))
	flag.BoolVar(&DisablePocScan, "nopoc", false, i18n.GetText("flag_no_poc"))

	// ═════════════════════════════════════════════════
	// Redis利用参数
	// ═════════════════════════════════════════════════
	flag.StringVar(&RedisFile, "rf", "", i18n.GetText("flag_redis_file"))
	flag.StringVar(&RedisShell, "rs", "", i18n.GetText("flag_redis_shell"))
	flag.BoolVar(&DisableRedis, "noredis", false, i18n.GetText("flag_disable_redis"))
	flag.StringVar(&RedisWritePath, "rwp", "", i18n.GetText("flag_redis_write_path"))
	flag.StringVar(&RedisWriteContent, "rwc", "", i18n.GetText("flag_redis_write_content"))
	flag.StringVar(&RedisWriteFile, "rwf", "", i18n.GetText("flag_redis_write_file"))

	// ═════════════════════════════════════════════════
	// 暴力破解控制参数
	// ═════════════════════════════════════════════════
	flag.BoolVar(&DisableBrute, "nobr", false, i18n.GetText("flag_disable_brute"))
	flag.IntVar(&MaxRetries, "retry", 3, i18n.GetText("flag_max_retries"))

	// ═════════════════════════════════════════════════
	// 输出与显示控制参数
	// ═════════════════════════════════════════════════
	flag.StringVar(&Outputfile, "o", "result.txt", i18n.GetText("flag_output_file"))
	flag.StringVar(&OutputFormat, "f", "txt", i18n.GetText("flag_output_format"))
	flag.BoolVar(&DisableSave, "no", false, i18n.GetText("flag_disable_save"))
	flag.BoolVar(&Silent, "silent", false, i18n.GetText("flag_silent_mode"))
	flag.BoolVar(&NoColor, "nocolor", false, i18n.GetText("flag_no_color"))
	flag.StringVar(&LogLevel, "log", LogLevelBaseInfoSuccess, i18n.GetText("flag_log_level"))
	flag.BoolVar(&DisableProgress, "nopg", false, i18n.GetText("flag_disable_progress"))

	// ═════════════════════════════════════════════════
	// 其他参数
	// ═════════════════════════════════════════════════
	flag.StringVar(&Shellcode, "sc", "", i18n.GetText("flag_shellcode"))
	flag.StringVar(&Language, "lang", "zh", i18n.GetText("flag_language"))

	// 帮助参数
	var showHelp bool
	flag.BoolVar(&showHelp, "help", false, i18n.GetText("flag_help"))

	// 解析命令行参数
	parseCommandLineArgs()

	// 设置语言
	i18n.SetLanguage(Language)


	// 更新进度条显示状态
	ShowProgress = !DisableProgress
	
	// 如果显示帮助或者没有提供目标，显示帮助信息并退出
	if showHelp || shouldShowHelp(Info) {
		flag.Usage()
		os.Exit(0)
	}
}

// parseCommandLineArgs 处理来自环境变量和命令行的参数
func parseCommandLineArgs() {
	// 首先检查环境变量中的参数
	envArgsString := os.Getenv("FS_ARGS")
	if envArgsString != "" {
		// 解析环境变量参数 (跨平台支持)
		envArgs, err := parseEnvironmentArgs(envArgsString)
		if err == nil && len(envArgs) > 0 {
			flag.CommandLine.Parse(envArgs)
			os.Unsetenv("FS_ARGS") // 使用后清除环境变量
			return
		}
		// 如果环境变量解析失败，继续使用命令行参数
	}

	// 解析命令行参数
	flag.Parse()
}

// parseEnvironmentArgs 安全地解析环境变量中的参数
func parseEnvironmentArgs(argsString string) ([]string, error) {
	if strings.TrimSpace(argsString) == "" {
		return nil, fmt.Errorf("empty arguments string")
	}

	// 使用更安全的参数分割方法
	var args []string
	var currentArg strings.Builder
	inQuote := false
	quoteChar := ' '

	for _, char := range argsString {
		switch {
		case char == '"' || char == '\'':
			if inQuote && char == quoteChar {
				inQuote = false
			} else if !inQuote {
				inQuote = true
				quoteChar = char
			} else {
				currentArg.WriteRune(char)
			}
		case char == ' ' && !inQuote:
			if currentArg.Len() > 0 {
				args = append(args, currentArg.String())
				currentArg.Reset()
			}
		default:
			currentArg.WriteRune(char)
		}
	}

	if currentArg.Len() > 0 {
		args = append(args, currentArg.String())
	}

	return args, nil
}

// preProcessLanguage 预处理语言参数，在定义flag之前设置语言
func preProcessLanguage() {
	// 遍历命令行参数查找-lang参数
	for i, arg := range os.Args {
		if arg == "-lang" && i+1 < len(os.Args) {
			lang := os.Args[i+1]
			if lang == "en" || lang == "zh" {
				Language = lang
				i18n.SetLanguage(lang)
				return
			}
		} else if strings.HasPrefix(arg, "-lang=") {
			lang := strings.TrimPrefix(arg, "-lang=")
			if lang == "en" || lang == "zh" {
				Language = lang
				i18n.SetLanguage(lang)
				return
			}
		}
	}
	
	// 检查环境变量
	envLang := os.Getenv("FS_LANG")
	if envLang == "en" || envLang == "zh" {
		Language = envLang
		i18n.SetLanguage(envLang)
	}
}

// shouldShowHelp 检查是否应该显示帮助信息
func shouldShowHelp(Info *HostInfo) bool {
	// 检查是否提供了扫描目标
	hasTarget := Info.Host != "" || TargetURL != "" || LocalMode
	
	// 如果没有提供任何扫描目标，则显示帮助
	if !hasTarget {
		return true
	}
	
	return false
}