fix: harden scan edge cases

This commit is contained in:
ZacharyZcR
2026-06-01 03:32:13 +08:00
parent 8b558b4f12
commit 8ec96bfe6d
21 changed files with 534 additions and 67 deletions
+33 -19
View File
@@ -49,7 +49,10 @@ func BuildConfig(fv *FlagVars, info *HostInfo) (*Config, *State, error) {
func parseCredentials(fv *FlagVars, cfg *Config) error {
// 解析用户名
usernames := parseUsernames(fv)
usernames, err := parseUsernames(fv)
if err != nil {
return err
}
if len(usernames) > 0 {
for serviceName := range cfg.Credentials.Userdict {
cfg.Credentials.Userdict[serviceName] = usernames
@@ -57,7 +60,10 @@ func parseCredentials(fv *FlagVars, cfg *Config) error {
}
// 解析密码
passwords := parsePasswords(fv)
passwords, err := parsePasswords(fv)
if err != nil {
return err
}
if len(passwords) > 0 {
cfg.Credentials.Passwords = passwords
}
@@ -84,7 +90,7 @@ func parseCredentials(fv *FlagVars, cfg *Config) error {
return nil
}
func parseUsernames(fv *FlagVars) []string {
func parseUsernames(fv *FlagVars) ([]string, error) {
var usernames []string
// 命令行用户名
@@ -102,7 +108,7 @@ func parseUsernames(fv *FlagVars) []string {
if lines, err := parsers.ReadLinesFromFile(fv.UsersFile); err == nil {
usernames = append(usernames, lines...)
} else {
LogError(i18n.Tr("config_read_users_failed", fv.UsersFile, err))
return nil, fmt.Errorf("%s", i18n.Tr("config_read_users_failed", fv.UsersFile, err))
}
}
@@ -116,15 +122,15 @@ func parseUsernames(fv *FlagVars) []string {
}
}
return removeDuplicate(usernames)
return removeDuplicate(usernames), nil
}
func parsePasswords(fv *FlagVars) []string {
func parsePasswords(fv *FlagVars) ([]string, error) {
var passwords []string
// 命令行密码
if fv.Password != "" {
passwords = append(passwords, splitCredentialValues(fv.Password)...)
passwords = append(passwords, fv.Password)
}
// 从文件读取
@@ -132,7 +138,7 @@ func parsePasswords(fv *FlagVars) []string {
if lines, err := parsers.ReadLinesFromFile(fv.PasswordsFile); err == nil {
passwords = append(passwords, lines...)
} else {
LogError(i18n.Tr("config_read_passwords_failed", fv.PasswordsFile, err))
return nil, fmt.Errorf("%s", i18n.Tr("config_read_passwords_failed", fv.PasswordsFile, err))
}
}
@@ -141,7 +147,7 @@ func parsePasswords(fv *FlagVars) []string {
passwords = append(passwords, splitCredentialValues(fv.AddPasswords)...)
}
return removeDuplicate(passwords)
return removeDuplicate(passwords), nil
}
func splitCredentialValues(input string) []string {
@@ -192,12 +198,15 @@ func parseHashes(fv *FlagVars) ([]string, [][]byte, error) {
// 命令行哈希
if fv.HashValue != "" {
hash := strings.TrimSpace(fv.HashValue)
if len(hash) == 32 {
hashValues = append(hashValues, hash)
if hashByte, err := hex.DecodeString(hash); err == nil {
hashBytes = append(hashBytes, hashByte)
}
if len(hash) != 32 {
return nil, nil, fmt.Errorf("invalid hash length: %s", hash)
}
hashByte, err := hex.DecodeString(hash)
if err != nil {
return nil, nil, err
}
hashValues = append(hashValues, hash)
hashBytes = append(hashBytes, hashByte)
}
// 从文件读取
@@ -225,13 +234,17 @@ func parseTargets(fv *FlagVars, info *HostInfo, cfg *Config, state *State) error
if port, portErr := strconv.Atoi(portStr); portErr == nil && port >= 1 && port <= 65535 {
// 有效的 host:port 格式
state.SetHostPorts([]string{info.Host})
info.Host = ""
ports = "" // 清空端口,避免双重扫描
}
}
}
// 解析 URL
urls := parseURLs(fv)
urls, err := parseURLs(fv)
if err != nil {
return err
}
if len(urls) > 0 {
state.SetURLs(urls)
if info.URL == "" && len(urls) == 1 {
@@ -247,7 +260,7 @@ func parseTargets(fv *FlagVars, info *HostInfo, cfg *Config, state *State) error
return nil
}
func parseURLs(fv *FlagVars) []string {
func parseURLs(fv *FlagVars) ([]string, error) {
var urls []string
// 命令行 URL
@@ -267,11 +280,11 @@ func parseURLs(fv *FlagVars) []string {
urls = append(urls, normalizeURL(line))
}
} else {
LogError(i18n.Tr("config_read_urls_failed", fv.URLsFile, err))
return nil, fmt.Errorf("%s", i18n.Tr("config_read_urls_failed", fv.URLsFile, err))
}
}
return removeDuplicate(urls)
return removeDuplicate(urls), nil
}
func normalizeURL(rawURL string) string {
@@ -279,7 +292,8 @@ func normalizeURL(rawURL string) string {
if rawURL == "" {
return rawURL
}
if !strings.HasPrefix(rawURL, "http://") && !strings.HasPrefix(rawURL, "https://") {
lowerURL := strings.ToLower(rawURL)
if !strings.HasPrefix(lowerURL, "http://") && !strings.HasPrefix(lowerURL, "https://") {
return "http://" + rawURL
}
return rawURL
+60 -4
View File
@@ -5,15 +5,71 @@ import (
"testing"
)
func TestParsePasswordsSplitsCommaAndWhitespace(t *testing.T) {
func TestParsePasswordsKeepsPrimaryPasswordLiteral(t *testing.T) {
fv := &FlagVars{
Password: "root,admin",
Password: "root admin",
AddPasswords: "pass1 pass2,pass3\tpass4",
}
got := parsePasswords(fv)
want := []string{"root", "admin", "pass1", "pass2", "pass3", "pass4"}
got, err := parsePasswords(fv)
if err != nil {
t.Fatalf("parsePasswords error = %v", err)
}
want := []string{"root admin", "pass1", "pass2", "pass3", "pass4"}
if !reflect.DeepEqual(got, want) {
t.Fatalf("parsePasswords() = %#v, want %#v", got, want)
}
}
func TestBuildConfigReturnsUserFileError(t *testing.T) {
_, _, err := BuildConfig(&FlagVars{UsersFile: "missing-users-file.txt"}, &HostInfo{})
if err == nil {
t.Fatal("BuildConfig should fail for missing users file")
}
}
func TestBuildConfigReturnsPasswordFileError(t *testing.T) {
_, _, err := BuildConfig(&FlagVars{PasswordsFile: "missing-passwords-file.txt"}, &HostInfo{})
if err == nil {
t.Fatal("BuildConfig should fail for missing passwords file")
}
}
func TestBuildConfigReturnsURLFileError(t *testing.T) {
_, _, err := BuildConfig(&FlagVars{URLsFile: "missing-urls-file.txt"}, &HostInfo{})
if err == nil {
t.Fatal("BuildConfig should fail for missing urls file")
}
}
func TestBuildConfigRejectsInvalidHashValue(t *testing.T) {
_, _, err := BuildConfig(&FlagVars{HashValue: "not-md5"}, &HostInfo{})
if err == nil {
t.Fatal("BuildConfig should fail for invalid hash value")
}
}
func TestParseTargetsHostPortDoesNotLeaveSyntheticHost(t *testing.T) {
fv := &FlagVars{Ports: "22"}
info := &HostInfo{Host: "127.0.0.1:8080"}
cfg := BuildConfigFromFlags(fv)
state := NewState()
if err := parseTargets(fv, info, cfg, state); err != nil {
t.Fatalf("parseTargets error = %v", err)
}
if info.Host != "" {
t.Fatalf("info.Host = %q, want empty after host:port extraction", info.Host)
}
if got := state.GetHostPorts(); !reflect.DeepEqual(got, []string{"127.0.0.1:8080"}) {
t.Fatalf("hostPorts = %#v, want host:port target", got)
}
}
func TestNormalizeURLKeepsUppercaseScheme(t *testing.T) {
got := normalizeURL("HTTPS://example.com")
if got != "HTTPS://example.com" {
t.Fatalf("normalizeURL() = %q", got)
}
}
+9 -4
View File
@@ -39,8 +39,11 @@ func NewHostIterator(host string, filename string, nohosts ...string) (*HostIter
sources = append(sources, hostSources...)
matcher := newHostMatcher()
if len(nohosts) > 0 && strings.TrimSpace(nohosts[0]) != "" {
if err := matcher.add(nohosts[0]); err != nil {
for _, exclude := range nohosts {
if strings.TrimSpace(exclude) == "" {
continue
}
if err := matcher.add(exclude); err != nil {
closeHostSources(sources)
return nil, err
}
@@ -180,10 +183,12 @@ func newFileHostSource(filename string) (*fileHostSource, error) {
if err != nil {
return nil, err
}
return &fileHostSource{
src := &fileHostSource{
file: file,
scanner: bufio.NewScanner(file),
}, nil
}
src.scanner.Buffer(make([]byte, 64*1024), 4*1024*1024)
return src, nil
}
func (s *fileHostSource) Next() (string, bool, error) {
+44
View File
@@ -2,7 +2,9 @@ package parsers
import (
"context"
"os"
"reflect"
"strings"
"testing"
)
@@ -59,3 +61,45 @@ func TestHostIteratorExcludeCIDR(t *testing.T) {
t.Fatalf("batch = %#v, want %#v", batch, want)
}
}
func TestHostIteratorAcceptsMultipleExcludeSources(t *testing.T) {
iter, err := NewHostIterator("192.168.1.0/29", "", "192.168.1.2", "192.168.1.5")
if err != nil {
t.Fatalf("NewHostIterator error = %v", err)
}
defer iter.Close()
batch, err := iter.NextBatch(context.Background(), 10)
if err != nil {
t.Fatalf("NextBatch error = %v", err)
}
want := []string{"192.168.1.1", "192.168.1.3", "192.168.1.4", "192.168.1.6"}
if !reflect.DeepEqual(batch, want) {
t.Fatalf("batch = %#v, want %#v", batch, want)
}
}
func TestHostIteratorReadsLongHostFileLine(t *testing.T) {
dir := t.TempDir()
path := dir + "/hosts.txt"
longPrefix := strings.Repeat("a", 70*1024)
host := longPrefix + ".example.com"
if err := os.WriteFile(path, []byte(host+"\n"), 0o600); err != nil {
t.Fatalf("WriteFile error = %v", err)
}
iter, err := NewHostIterator("", path)
if err != nil {
t.Fatalf("NewHostIterator error = %v", err)
}
defer iter.Close()
batch, err := iter.NextBatch(context.Background(), 1)
if err != nil {
t.Fatalf("NextBatch error = %v", err)
}
if !reflect.DeepEqual(batch, []string{host}) {
t.Fatalf("batch = %#v, want long host", batch)
}
}
+16 -2
View File
@@ -696,6 +696,18 @@ func TestParseIP_Exclude(t *testing.T) {
}
}
func TestParseIPMultipleExcludeSources(t *testing.T) {
result, err := ParseIP("192.168.1.1-192.168.1.4", "", "192.168.1.2", "192.168.1.4")
if err != nil {
t.Fatalf("ParseIP error = %v", err)
}
expected := []string{"192.168.1.1", "192.168.1.3"}
if !reflect.DeepEqual(result, expected) {
t.Fatalf("ParseIP with multiple excludes = %v, want %v", result, expected)
}
}
// TestParseIP_Deduplicate 测试去重
func TestParseIP_Deduplicate(t *testing.T) {
result, err := ParseIP("192.168.1.1,192.168.1.1,192.168.1.2,192.168.1.2", "", "")
@@ -800,7 +812,9 @@ func TestParsePortRange(t *testing.T) {
// TestExcludeHosts 测试排除主机
func TestExcludeHosts(t *testing.T) {
hosts := []string{"host1", "host2", "host3", "host4"}
exclude := []string{"host2", "host4"}
exclude := newHostMatcher()
exclude.exact["host2"] = struct{}{}
exclude.exact["host4"] = struct{}{}
result := excludeFromList(hosts, exclude)
expected := []string{"host1", "host3"}
@@ -815,7 +829,7 @@ func TestExcludeHosts(t *testing.T) {
// TestExcludeHosts_EmptyExclude 测试空排除列表
func TestExcludeHosts_EmptyExclude(t *testing.T) {
hosts := []string{"host1", "host2"}
result := excludeFromList(hosts, []string{})
result := excludeFromList(hosts, nil)
if !reflect.DeepEqual(result, hosts) {
t.Errorf("excludeFromList(空排除列表) 应该返回原列表")
+19 -14
View File
@@ -53,18 +53,27 @@ func ParseIP(host string, filename string, nohosts ...string) ([]string, error)
if host != "" {
hostList, err := parseHostString(host)
if err != nil {
return nil, fmt.Errorf(i18n.GetText("parser_parse_host_failed")+": %w", err)
return nil, fmt.Errorf(i18n.GetText("parser_parse_host_failed")+": %w", err)
}
hosts = append(hosts, hostList...)
}
// 处理排除主机
if len(nohosts) > 0 && nohosts[0] != "" {
excludeList, err := parseHostString(nohosts[0])
if err != nil {
return nil, fmt.Errorf(i18n.GetText("parser_parse_exclude_failed")+": %w", err)
if len(nohosts) > 0 {
matcher := newHostMatcher()
hasExclude := false
for _, exclude := range nohosts {
if strings.TrimSpace(exclude) == "" {
continue
}
hasExclude = true
if err := matcher.add(exclude); err != nil {
return nil, fmt.Errorf(i18n.GetText("parser_parse_exclude_failed")+": %w", err)
}
}
if hasExclude {
hosts = excludeFromList(hosts, matcher)
}
hosts = excludeFromList(hosts, excludeList)
}
// 去重和排序
@@ -212,6 +221,7 @@ func ReadLinesFromFile(filename string) ([]string, error) {
var lines []string
scanner := bufio.NewScanner(file)
scanner.Buffer(make([]byte, 64*1024), 4*1024*1024)
for scanner.Scan() {
line := strings.TrimSpace(scanner.Text())
if line != "" && !strings.HasPrefix(line, "#") {
@@ -420,19 +430,14 @@ func incrementIP(ip net.IP) {
}
// excludeFromList 从列表中排除指定项
func excludeFromList(hosts, excludeList []string) []string {
if len(excludeList) == 0 {
func excludeFromList(hosts []string, matcher *hostMatcher) []string {
if matcher == nil {
return hosts
}
excludeMap := make(map[string]struct{}, len(excludeList))
for _, e := range excludeList {
excludeMap[e] = struct{}{}
}
result := make([]string, 0, len(hosts))
for _, h := range hosts {
if _, found := excludeMap[h]; !found {
if !matcher.match(h) {
result = append(result, h)
}
}