package main

import (
	"bufio"
	"bytes"
	"database/sql"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"math/rand"
	"net"
	"net/http"
	"os"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	_ "github.com/go-sql-driver/mysql"
	_ "github.com/mattn/go-sqlite3"
	"golang.org/x/net/html"
)

type DBConfig struct {
	DbDriver string `json:"dbdriver"`
	Username string `json:"username"`
	Password string `json:"password"`
	DBName   string `json:"dbname"`
}

type Webconfig struct {
	Port     int    `json:"port"`
	Username string `json:"username"`
	Password string `json:"password"`
}
type Config struct {
	ZmapFile       string    `json:"zmapfile"`
	Channelfeeders int       `json:"channelfeeders"`
	DB             DBConfig  `json:"db"`
	Web            Webconfig `json:"web"`
}

func loadConfig(path string) (*Config, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	var cfg Config
	if err := json.NewDecoder(f).Decode(&cfg); err != nil {
		return nil, err
	}
	return &cfg, nil
}

type BannerStore interface {
	InsertHTTPBanner(ip string, listening bool, title, resp, errStr, headers, serverHeader, contentType, contentLength, redirect, metaTags, bodySnippet string) error
	InsertSSHBanner(ip string, listening bool, firstline string) error
	InsertTelnetBanner(ip string, listening bool, firstline, fullbanner string) error
}
type SQLBannerStore struct {
	db           *sql.DB
	insertHTTP   *sql.Stmt
	insertSSH    *sql.Stmt
	insertTelnet *sql.Stmt
}

func (s *SQLBannerStore) InsertHTTPBanner(
	ip string,
	listening bool,
	title, resp, errStr, headers,
	serverHeader, contentType, contentLength, redirect, metaTags, bodySnippet string,
) error {
	_, err := s.insertHTTP.Exec(ip, listening, title, resp, errStr, headers,
		serverHeader, contentType, contentLength, redirect, metaTags, bodySnippet)
	return err
}

func (s *SQLBannerStore) InsertSSHBanner(ip string, listening bool, firstline string) error {
	_, err := s.insertSSH.Exec(ip, listening, firstline)
	return err
}

func (s *SQLBannerStore) InsertTelnetBanner(ip string, listening bool, firstline, fullbanner string) error {
	_, err := s.insertTelnet.Exec(ip, listening, firstline, fullbanner)
	return err
}

var (
	db          *sql.DB
	defaultFile = "zmaplist.txt"
	store       BannerStore
	excluded    []ipRange
	allowed     []ipRange
)

const red = "\033[31m"
const reset = "\033[0m"
const green = "\033[32m"
const blue = "\033[34m"

type ipRange struct {
	start uint32
	end   uint32
}

func ipToUint32(ip net.IP) uint32 {
	ip = ip.To4()
	return uint32(ip[0])<<24 | uint32(ip[1])<<16 | uint32(ip[2])<<8 | uint32(ip[3])
}

func uint32ToIP(n uint32) net.IP {
	return net.IPv4(
		byte(n>>24),
		byte(n>>16),
		byte(n>>8),
		byte(n),
	)
}
func parseCIDRRange(cidr string) ipRange {
	_, network, _ := net.ParseCIDR(cidr)
	start := ipToUint32(network.IP)
	mask := ipToUint32(net.IP(network.Mask))
	hostBits := ^mask
	end := start | hostBits
	return ipRange{start, end}
}
func main() {
	blocks := []string{
		"10.0.0.0/8", "172.16.0.0/12", "192.168.0.0/16",
		"127.0.0.0/8", "169.254.0.0/16", "100.64.0.0/10",
		"192.0.0.0/24", "192.0.2.0/24", "192.88.99.0/24",
		"198.18.0.0/15", "198.51.100.0/24", "203.0.113.0/24",
		"224.0.0.0/4", "240.0.0.0/4", "255.255.255.255/32",
		// US Gov / DoD legacy
		"6.0.0.0/8", "7.0.0.0/8", "11.0.0.0/8",
		"21.0.0.0/8", "22.0.0.0/8", "26.0.0.0/8",
		"28.0.0.0/8", "29.0.0.0/8", "30.0.0.0/8",
		"33.0.0.0/8", "55.0.0.0/8",
	}
	for _, b := range blocks {
		excluded = append(excluded, parseCIDRRange(b))
	}
	for first := 1; first <= 255; first++ {
		cidr := fmt.Sprintf("%d.0.0.0/8", first)
		r := parseCIDRRange(cidr)
		skip := false
		for _, ex := range excluded {
			if r.start >= ex.start && r.start <= ex.end {
				skip = true
				break
			}
		}
		if !skip {
			allowed = append(allowed, r)
		}
	}

	rand.Seed(time.Now().UnixNano())
	cfg, err := loadConfig("config.json")
	if err != nil {
		log.Fatal(err)
	}

	switch cfg.DB.DbDriver {
	case "mysql":
		store, err = sqlInitMySQL(fmt.Sprintf("%s:%s@tcp(127.0.0.1:3306)/%s", cfg.DB.Username, cfg.DB.Password, cfg.DB.DBName))
	case "sqlite":
		store, err = sqlInitSQLite(cfg.DB.DBName)
	default:
		log.Fatalf("Unsupported driver: %s", cfg.DB.DbDriver)
	}

	if err != nil {
		log.Fatal(err)
	}
	fmt.Println("Starting IP scan")
	ConcurrentScan(defaultFile, cfg)
	fmt.Println("Finished IP scan!")
	return
}

func sqlInitMySQL(connargs string) (*SQLBannerStore, error) {
	// open connection
	var err error
	db, err = sql.Open("mysql", connargs)
	if err != nil {
		return nil, err
	}

	// create tables
	tables := []string{
		`CREATE TABLE IF NOT EXISTS httpbanners (
        Id INT AUTO_INCREMENT PRIMARY KEY,
        Ip VARCHAR(255) NOT NULL UNIQUE,
        Scanned_on TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        IsListening TINYINT(1) NOT NULL,
        httptitle TEXT,
        Httpresp TEXT,
        Headers TEXT,
        Error TEXT,
        ServerHeader TEXT,
        ContentType TEXT,
        ContentLength TEXT,
        Redirect TEXT,
        MetaTags TEXT,
        BodySnippet TEXT
    )`,

		`CREATE TABLE IF NOT EXISTS sshbanners (
        Id INT AUTO_INCREMENT PRIMARY KEY,
        Ip VARCHAR(255) NOT NULL UNIQUE,
        scanned_on TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        IsListening TINYINT(1) NOT NULL,
        firstline TEXT
    )`,

		`CREATE TABLE IF NOT EXISTS telnetbanners (
        Id INT AUTO_INCREMENT PRIMARY KEY,
        Ip VARCHAR(255) NOT NULL UNIQUE,
        Scanned_on TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        IsListening TINYINT(1) NOT NULL,
        firstline TEXT,
        fullbanner TEXT
    )`,
	}

	for _, stmt := range tables {
		if _, err := db.Exec(stmt); err != nil {
			return nil, err
		}
	}

	// prepare statements
	httpStmt, err := db.Prepare(`
    INSERT INTO httpbanners
        (Ip, IsListening, httptitle, Httpresp, Error, Headers,
         ServerHeader, ContentType, ContentLength, Redirect, MetaTags, BodySnippet)
    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
    ON DUPLICATE KEY UPDATE
        IsListening=VALUES(IsListening),
        httptitle=VALUES(httptitle),
        Httpresp=VALUES(Httpresp),
        Error=VALUES(Error),
        Headers=VALUES(Headers),
        ServerHeader=VALUES(ServerHeader),
        ContentType=VALUES(ContentType),
        ContentLength=VALUES(ContentLength),
        Redirect=VALUES(Redirect),
        MetaTags=VALUES(MetaTags),
        BodySnippet=VALUES(BodySnippet)
`)

	if err != nil {
		return nil, err
	}

	sshStmt, err := db.Prepare(`
        INSERT INTO sshbanners (Ip, IsListening, firstline)
        VALUES (?, ?, ?)
        ON DUPLICATE KEY UPDATE
            IsListening=VALUES(IsListening),
            firstline=VALUES(firstline)
    `)
	if err != nil {
		return nil, err
	}

	telnetStmt, err := db.Prepare(`
        INSERT INTO telnetbanners (Ip, IsListening, firstline, fullbanner)
        VALUES (?, ?, ?, ?)
        ON DUPLICATE KEY UPDATE
            IsListening=VALUES(IsListening),
            firstline=VALUES(firstline),
            fullbanner=VALUES(fullbanner)
    `)
	if err != nil {
		return nil, err
	}

	return &SQLBannerStore{
		db:           db,
		insertHTTP:   httpStmt,
		insertSSH:    sshStmt,
		insertTelnet: telnetStmt,
	}, nil
}

func sqlInitSQLite(dbname string) (*SQLBannerStore, error) {
	db, err := sql.Open("sqlite3", dbname)
	if err != nil {
		return nil, err
	}
	tables := []string{
		`CREATE TABLE IF NOT EXISTS httpbanners (
            Id INTEGER PRIMARY KEY AUTOINCREMENT,
            Ip TEXT NOT NULL UNIQUE,
            Scanned_on TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            IsListening INTEGER NOT NULL,
            httptitle TEXT,
            Httpresp TEXT,
            Headers TEXT,
            Error TEXT,
            ServerHeader TEXT,
            ContentType TEXT,
            ContentLength TEXT,
            Redirect TEXT,
            MetaTags TEXT,
            BodySnippet TEXT
        )`,
		`CREATE TABLE IF NOT EXISTS sshbanners (
            Id INTEGER PRIMARY KEY AUTOINCREMENT,
            Ip TEXT NOT NULL UNIQUE,
            Scanned_on TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            IsListening INTEGER NOT NULL,
            firstline TEXT
        )`,
		`CREATE TABLE IF NOT EXISTS telnetbanners (
            Id INTEGER PRIMARY KEY AUTOINCREMENT,
            Ip TEXT NOT NULL UNIQUE,
            Scanned_on TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            IsListening INTEGER NOT NULL,
            firstline TEXT,
            fullbanner TEXT
        )`,
	}

	for _, stmt := range tables {
		if _, err := db.Exec(stmt); err != nil {
			return nil, err
		}
	}
	httpStmt, err := db.Prepare(`
        INSERT INTO httpbanners
            (Ip, IsListening, httptitle, Httpresp, Error, Headers,
             ServerHeader, ContentType, ContentLength, Redirect, MetaTags, BodySnippet)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        ON CONFLICT(Ip) DO UPDATE SET
            IsListening=excluded.IsListening,
            httptitle=excluded.httptitle,
            Httpresp=excluded.Httpresp,
            Error=excluded.Error,
            Headers=excluded.Headers,
            ServerHeader=excluded.ServerHeader,
            ContentType=excluded.ContentType,
            ContentLength=excluded.ContentLength,
            Redirect=excluded.Redirect,
            MetaTags=excluded.MetaTags,
            BodySnippet=excluded.BodySnippet
    `)
	if err != nil {
		return nil, err
	}

	sshStmt, err := db.Prepare(`
        INSERT INTO sshbanners (Ip, IsListening, firstline)
        VALUES (?, ?, ?)
        ON CONFLICT(Ip) DO UPDATE SET
            IsListening=excluded.IsListening,
            firstline=excluded.firstline
    `)
	if err != nil {
		return nil, err
	}

	telnetStmt, err := db.Prepare(`
        INSERT INTO telnetbanners (Ip, IsListening, firstline, fullbanner)
        VALUES (?, ?, ?, ?)
        ON CONFLICT(Ip) DO UPDATE SET
            IsListening=excluded.IsListening,
            firstline=excluded.firstline,
            fullbanner=excluded.fullbanner
    `)
	if err != nil {
		return nil, err
	}

	return &SQLBannerStore{
		db:           db,
		insertHTTP:   httpStmt,
		insertSSH:    sshStmt,
		insertTelnet: telnetStmt,
	}, nil
}
func isPrivateIP(ip net.IP) bool { // blacklisted CIDR ranges from being scanned, you can also put your own IPs, a long as its in CIDR notation
	var privateBlocks = []string{
		"10.0.0.0/8",         // RFC1918 private
		"172.16.0.0/12",      // RFC1918 private
		"192.168.0.0/16",     // RFC1918 private
		"127.0.0.0/8",        // loopback
		"169.254.0.0/16",     // link-local
		"100.64.0.0/10",      // carrier-grade NAT
		"192.0.0.0/24",       // IANA special use
		"192.0.2.0/24",       // TEST-NET-1
		"192.88.99.0/24",     // deprecated 6to4 relay
		"198.18.0.0/15",      // benchmarking
		"198.51.100.0/24",    // TEST-NET-2
		"203.0.113.0/24",     // TEST-NET-3
		"224.0.0.0/4",        // multicast
		"240.0.0.0/4",        // reserved
		"255.255.255.255/32", // broadcast
	}

	// US Government / DoD legacy ranges (policy exclusions)
	var govBlocks = []string{
		"6.0.0.0/8",
		"7.0.0.0/8",
		"11.0.0.0/8",
		"21.0.0.0/8",
		"22.0.0.0/8",
		"26.0.0.0/8",
		"28.0.0.0/8",
		"29.0.0.0/8",
		"30.0.0.0/8",
		"33.0.0.0/8",
		"55.0.0.0/8",
	}
	for _, block := range append(privateBlocks, govBlocks...) {
		_, cidr, _ := net.ParseCIDR(block)
		if cidr.Contains(ip) {
			return true
		}
	}
	return false
}
func generateRandomPublicIP() string {
	r := allowed[rand.Intn(len(allowed))]
	ip := r.start + uint32(rand.Intn(int(r.end-r.start+1)))
	return uint32ToIP(ip).String()
}
func grabbanner(ip string) {
	ports := []int{80, 8080, 8088}
	client := http.Client{Timeout: 1500 * time.Millisecond}

	for _, port := range ports {
		url := fmt.Sprintf("http://%s:%d", ip, port)
		resp, err := client.Get(url)
		if err != nil {
			//store.InsertHTTPBanner(fmt.Sprintf("%s:%d", ip, port),false, "", "", err.Error(), "", "", "", "", "", "", "")
			continue
		}
		defer resp.Body.Close()

		// Limit body size to avoid huge pages
		body, err := io.ReadAll(io.LimitReader(resp.Body, 8192))
		if err != nil {
			store.InsertHTTPBanner(fmt.Sprintf("%s:%d", ip, port), true, "", "", err.Error(), "", "", "", "", "", "", "")
			continue
		}

		htmlStr := string(body)
		title := extractTitle(htmlStr)
		metas := extractMeta(htmlStr)

		// Collect headers
		var sb strings.Builder
		for k, v := range resp.Header {
			sb.WriteString(fmt.Sprintf("%s: %s\n", k, strings.Join(v, ", ")))
		}
		headers := sb.String()

		// Extra metadata
		server := resp.Header.Get("Server")
		contentType := resp.Header.Get("Content-Type")
		contentLength := resp.Header.Get("Content-Length")
		redirect := resp.Header.Get("Location")

		// Truncate body snippet
		snippet := htmlStr
		if len(snippet) > 2000 {
			snippet = snippet[:2000]
		}

		fmt.Println(green+"[ HTTP BANNER ] target found", url, reset)
		store.InsertHTTPBanner(fmt.Sprintf("%s:%d", ip, port), true,
			title,
			fmt.Sprintf("%d %s", resp.StatusCode, resp.Status),
			"", headers,
			server, contentType, contentLength, redirect,
			fmt.Sprint(metas), snippet)
	}
}

func extractTitle(htmlStr string) string { // finally it works..except im trusting the go devs html package for this
	doc, err := html.Parse(strings.NewReader(htmlStr))
	if err != nil {
		return ""
	}
	var f func(*html.Node) string
	f = func(n *html.Node) string {
		if n.Type == html.ElementNode && n.Data == "title" && n.FirstChild != nil {
			return n.FirstChild.Data
		}
		for c := n.FirstChild; c != nil; c = c.NextSibling {
			if t := f(c); t != "" {
				return t
			}
		}
		return ""
	}
	return f(doc)
}
func extractMeta(htmlStr string) map[string]string {
	metas := make(map[string]string)
	doc, err := html.Parse(strings.NewReader(htmlStr))
	if err != nil {
		return metas
	}
	var f func(*html.Node)
	f = func(n *html.Node) {
		if n.Type == html.ElementNode && n.Data == "meta" {
			var name, content string
			for _, attr := range n.Attr {
				if attr.Key == "name" {
					name = attr.Val
				}
				if attr.Key == "content" {
					content = attr.Val
				}
			}
			if name != "" {
				metas[name] = content
			}
		}
		for c := n.FirstChild; c != nil; c = c.NextSibling {
			f(c)
		}
	}
	f(doc)
	return metas
}
func startGenerators(n int, httpJobs, sshJobs, telnetJobs chan string) {
	for i := 0; i < n; i++ {
		go func(seed int64) {
			rand.Seed(time.Now().UnixNano() + seed)
			for {
				ip := generateRandomPublicIP()
				select {
				case httpJobs <- ip:
				default:
				}
				select {
				case sshJobs <- ip:
				default:
				}
				select {
				case telnetJobs <- ip:
				default:
					// meh..
				}
			}
		}(int64(i))
	}
}

func ConcurrentScan(filename string, conf *Config) { // the big main loop
	httpJobs := make(chan string, 5700)
	telnetJobs := make(chan string, 5000)
	sshJobs := make(chan string, 5000)

	var wg sync.WaitGroup

	var httpCount, sshCount, telnetCount uint64
	// helper to spawn workers
	spawn := func(n int, jobs chan string, fn func(string), counter *uint64) {
		for w := 0; w < n; w++ {
			wg.Add(1)
			go func() {
				defer wg.Done()
				for ip := range jobs {
					fn(ip)
					atomic.AddUint64(counter, 1)
				}
			}()
		}
	}

	// spawn workers
	spawn(340, httpJobs, grabbanner, &httpCount)
	spawn(150, sshJobs, func(ip string) { BannerSSH(fmt.Sprintf("%s:22", ip)) }, &sshCount)
	spawn(150, telnetJobs, BannerTelnet, &telnetCount)
	// feed ZMap file if present
	file, err := os.Open(filename)
	if err == nil {
		scanner := bufio.NewScanner(file)
		fed := 0
		for scanner.Scan() {
			ip := strings.TrimSpace(scanner.Text())
			if ip == "" {
				continue
			}
			httpJobs <- ip
			telnetJobs <- ip
			sshJobs <- ip
			fed++
		}
		file.Close()
		if fed > 0 {
			fmt.Printf("Fed %d IPs from ZMap file\n", fed)
		} else {
			fmt.Println("ZMap file contained no usable IPs")
		}
	} else {
		fmt.Println("No ZMap file found")
	}
	go startGenerators(conf.Channelfeeders, httpJobs, sshJobs, telnetJobs) // parallel feeders
	go func() {
		ticker := time.NewTicker(20 * time.Second)
		defer ticker.Stop()
		for range ticker.C {
			fmt.Printf("[STATUS] HTTP:%d SSH:%d TELNET:%d | Queues H:%d S:%d T:%d\n",
				atomic.LoadUint64(&httpCount),
				atomic.LoadUint64(&sshCount),
				atomic.LoadUint64(&telnetCount),
				len(httpJobs), len(sshJobs), len(telnetJobs))

		}
	}()
	wg.Wait()
}

func BannerTelnet(ip string) {
	conn, err := net.DialTimeout("tcp", fmt.Sprintf("%s:23", ip), 1*time.Second)
	if err != nil {
		// host not listening
		//	store.InsertTelnetBanner(ip, false, "", "")
		return
	}
	defer conn.Close()
	conn.SetReadDeadline(time.Now().Add(4 * time.Second))

	buf := make([]byte, 2048)
	var sb bytes.Buffer
	for {
		n, err := conn.Read(buf)
		if err != nil {
			break
		}
		sb.Write(buf[:n])
	}

	clean := stripTelnetControls(sb.Bytes())
	finalbanner := strings.TrimSpace(string(clean))
	firstLine := strings.SplitN(finalbanner, "\n", 2)[0]

	fmt.Println(blue+"[TELNET SCANNER] HOST FOUND", ip, reset)
	store.InsertTelnetBanner(ip, true, firstLine, finalbanner)
}

func BannerSSH(ip string) {
	conn, err := net.DialTimeout("tcp", ip, 1500*time.Millisecond)
	if err != nil {
		// host not listening
		//store.InsertSSHBanner(ip, false, "")
		return
	}
	defer conn.Close()

	reader := bufio.NewReader(conn)
	banner, err := reader.ReadString('\n')
	if err != nil {
		// connection established but banner read failed
		//store.InsertSSHBanner(ip, true, "")
		return
	}

	// success: got a banner line
	store.InsertSSHBanner(ip, true, banner)
}

// StripTelnetControls used to remove the binary data from telnet connections to extract strings
func stripTelnetControls(data []byte) []byte { // the ANNOYING IAC CODES
	out := make([]byte, 0, len(data))
	i := 0
	for i < len(data) {
		if data[i] == 0xFF && i+2 < len(data) {
			i += 3
			continue
		}
		out = append(out, data[i])
		i++
	}
	return out
}