better infinity mode implementation

Author: juzeon

main.go

@@ -66,7 +66,7 @@ func main() {
 	}
 	var hostChan <-chan Host
 	if addr != "" {
-		hostChan = Iterate(strings.NewReader(addr))
+		hostChan = IterateAddr(addr)
 	} else if in != "" {
 		f, err := os.Open(in)
 		if err != nil {

scanner.go

@@ -11,22 +11,12 @@ import (
 
 func ScanTLS(host Host, out chan<- string) {
 	if host.IP == nil {
-		ips, err := net.LookupIP(host.Origin)
+		ip, err := LookupIP(host.Origin)
 		if err != nil {
-			slog.Debug("Failed to lookup", "origin", host.Origin, "err", err)
+			slog.Debug("Failed to get IP from the origin", "origin", host.Origin, "err", err)
 			return
 		}
-		var arr []net.IP
-		for _, ip := range ips {
-			if ip.To4() != nil || enableIPv6 {
-				arr = append(arr, ip)
-			}
-		}
-		if len(arr) == 0 {
-			slog.Debug("No IP found", "origin", host.Origin)
-			return
-		}
-		host.IP = arr[0]
+		host.IP = ip
 	}
 	hostPort := net.JoinHostPort(host.IP.String(), strconv.Itoa(port))
 	conn, err := net.DialTimeout("tcp", hostPort, time.Duration(timeout)*time.Second)

utils.go

@@ -3,8 +3,11 @@ package main
 import (
 	"bufio"
 	"errors"
+	"fmt"
 	"io"
 	"log/slog"
+	"math"
+	"math/big"
 	"net"
 	"net/netip"
 	"regexp"
@@ -108,6 +111,67 @@ func ExistOnlyOne(arr []string) bool {
 	}
 	return exist
 }
+func IterateAddr(addr string) <-chan Host {
+	hostChan := make(chan Host)
+	_, _, err := net.ParseCIDR(addr)
+	if err == nil {
+		// is CIDR
+		return Iterate(strings.NewReader(addr))
+	}
+	ip := net.ParseIP(addr)
+	if ip == nil {
+		ip, err = LookupIP(addr)
+		if err != nil {
+			close(hostChan)
+			slog.Error("Not a valid IP, IP CIDR or domain", "addr", addr)
+			return hostChan
+		}
+	}
+	go func() {
+		slog.Info("Enable infinite mode", "init", ip.String())
+		lowIP := ip
+		highIP := ip
+		hostChan <- Host{
+			IP:     ip,
+			Origin: addr,
+			Type:   HostTypeIP,
+		}
+		for i := 0; i < math.MaxInt; i++ {
+			if i%2 == 0 {
+				lowIP = NextIP(lowIP, false)
+				hostChan <- Host{
+					IP:     lowIP,
+					Origin: lowIP.String(),
+					Type:   HostTypeIP,
+				}
+			} else {
+				highIP = NextIP(highIP, true)
+				hostChan <- Host{
+					IP:     highIP,
+					Origin: highIP.String(),
+					Type:   HostTypeIP,
+				}
+			}
+		}
+	}()
+	return hostChan
+}
+func LookupIP(addr string) (net.IP, error) {
+	ips, err := net.LookupIP(addr)
+	if err != nil {
+		return nil, fmt.Errorf("failed to lookup: %w", err)
+	}
+	var arr []net.IP
+	for _, ip := range ips {
+		if ip.To4() != nil || enableIPv6 {
+			arr = append(arr, ip)
+		}
+	}
+	if len(arr) == 0 {
+		return nil, errors.New("no IP found")
+	}
+	return arr[0], nil
+}
 func RemoveDuplicateStr(strSlice []string) []string {
 	allKeys := make(map[string]bool)
 	var list []string
@@ -128,3 +192,17 @@ func OutWriter(writer io.Writer) chan<- string {
 	}()
 	return ch
 }
+func NextIP(ip net.IP, increment bool) net.IP {
+	// Convert to big.Int and increment
+	ipb := big.NewInt(0).SetBytes(ip)
+	if increment {
+		ipb.Add(ipb, big.NewInt(1))
+	} else {
+		ipb.Sub(ipb, big.NewInt(1))
+	}
+
+	// Add leading zeros
+	b := ipb.Bytes()
+	b = append(make([]byte, len(ip)-len(b)), b...)
+	return b
+}