targets.go

package engine

import (
	"fmt"
	"io"
	"net"
	"sync"
	"time"
)

type Target struct {
	IP   string
	Port uint16
}

func (t Target) String() string {
	return fmt.Sprintf("%s:%d", t.IP, t.Port)
}

var ips []OriginString
var ports []OriginUint16

var targets struct {
	sync.RWMutex
	targets []Target
}

var isProxyEnvironment bool

func AddIP(ip string, origin interface{}) {
	for _, exists := range ips {
		if exists.string == ip {
			return
		}
	}

	Logger.Debugf("Collecting ip: %v\n", ip)
	ips = append(ips, OriginString{Origin{origin}, ip})
}

func AddPort(port uint16, origin interface{}) {
	for _, exists := range ports {
		if exists.uint16 == port {
			return
		}
	}

	Logger.Debugf("Collecting port: %v\n", port)
	ports = append(ports, OriginUint16{Origin{origin}, port})
}

func GetIPsByOrigin(origin interface{}) []string {
	var res []string
	for _, ip := range ips {
		if ip.Origin.Origin == origin {
			res = append(res, ip.string)
		}
	}
	return res
}

func GetPortsByOrigin(origin interface{}) []uint16 {
	var res []uint16
	for _, port := range ports {
		if port.Origin.Origin == origin {
			res = append(res, port.uint16)
		}
	}
	return res
}

func connectToTarget(target Target) (net.Conn, error) {
	tcpAddr, err := net.ResolveTCPAddr("tcp4", target.String())
	if err != nil {
		return nil, err
	}

	conn, err := net.DialTimeout("tcp", tcpAddr.String(), time.Millisecond*50)
	if err != nil {
		Logger.Debugf("Target %v isn't reachable", target)
		return nil, err
	}

	return conn, nil
}

func isTargetReachable(target Target) bool {
	conn, err := connectToTarget(target)
	if err != nil {
		return false
	}
	defer conn.Close()

	if isProxyEnvironment {
		one := make([]byte, 1)
		conn.SetReadDeadline(time.Now().Add(2 * time.Second))
		if _, err := conn.Read(one); err != nil {
			if err == io.EOF {
				return false
			}
		}
	}

	return true
}

func checkIsProxyEnvironment() bool {
	if isTargetReachable(Target{"8.8.8.8", 1337}) {
		return true
	}
	return false
}

func targetChecker(wg *sync.WaitGroup, workerid int, gen <-chan Target) {
	defer wg.Done()

	Logger.Infof("[%d] Started.", workerid)

	for target := range gen {
		if isTargetReachable(target) {
			Logger.Infof("[%d] Target %s reachable!", workerid, target)

			targets.Lock()
			targets.targets = append(targets.targets, target)
			targets.Unlock()
		}
	}

}

func targetProducer() <-chan Target {
	ch := make(chan Target)
	go func() {
		for _, ip := range ips {
			for _, port := range ports {
				target := Target{ip.string, port.uint16}
				ch <- target
			}
		}
		close(ch)
	}()
	return ch
}

func CheckTargets(workers int) {
	Logger.Infof("Checking targets with %d workers", workers)
	isProxyEnvironment = checkIsProxyEnvironment()
	if isProxyEnvironment {
		Logger.Infof("Proxy environment detected!")
	}

	Logger.Infof("Starting %d workers", workers)
	producer := targetProducer()

	var wg sync.WaitGroup
	wg.Add(workers)
	for i := 0; i < workers; i++ {
		go targetChecker(&wg, i, producer)
	}
	wg.Wait()

	Logger.Infof("Done")
}

func GenerateTarget() <-chan Target {
	ch := make(chan Target)
	go func() {
		for _, target := range targets.targets {
			ch <- target
		}
		close(ch)
	}()
	return ch
}