star/net/monitorip.go

package net

import (
	"b612.me/apps/b612/netforward"
	"b612.me/starlog"
	"fmt"
	"net"
	"sync"
	"time"
)

type Monitor struct {
	IPs          []string
	Port         int
	ScanType     string
	Timeout      int
	Interval     int
	Log          string
	Retry        int
	Threads      int
	WithHostname bool
	pinger       *Pinger
	status       map[string]bool
	statusLock   sync.RWMutex
	stopChan     chan struct{}
}

func NewMonitor(ips []string, scanType string) *Monitor {
	return &Monitor{
		IPs:      ips,
		ScanType: scanType,
		status:   make(map[string]bool),
		stopChan: make(chan struct{}),
		Timeout:  1000,
		Interval: 1,
		Retry:    1,
		Threads:  50,
	}
}

func (m *Monitor) Start() error {
	if m.Log != "" {
		starlog.SetLogFile(m.Log, starlog.Std, true)
	}
	m.pinger = NewPinger(1127)
	// Initialize status
	m.statusLock.Lock()
	for _, ip := range m.IPs {
		m.status[ip] = false // Initial state as down
	}
	m.statusLock.Unlock()

	ticker := time.NewTicker(time.Duration(m.Interval) * time.Second)
	defer ticker.Stop()
	m.checkAllIPs()
	m.displayStatus()
	for {
		select {
		case <-ticker.C:
			m.checkAllIPs()
			m.displayStatus()
		case <-m.stopChan:
			return nil
		}
	}
}

func (m *Monitor) Stop() {
	close(m.stopChan)
}

func (m *Monitor) checkAllIPs() {
	var wg sync.WaitGroup
	sem := make(chan struct{}, m.Threads)

	for _, ip := range m.IPs {
		wg.Add(1)
		sem <- struct{}{}

		go func(ip string) {
			defer func() {
				<-sem
				wg.Done()
			}()

			currentStatus := m.checkIP(ip)
			m.updateStatus(ip, currentStatus)
		}(ip)
	}
	wg.Wait()
}

func (m *Monitor) checkIP(ip string) bool {
	for i := 0; i < m.Retry+1; i++ {
		var success bool
		var err error

		switch m.ScanType {
		case "icmp":
			err = m.pinger.Ping(ip, time.Duration(m.Timeout)*time.Millisecond)
			success = err == nil
		case "tcp":
			dialer := net.Dialer{
				Timeout: time.Duration(m.Timeout) * time.Millisecond,
				Control: netforward.ControlSetReUseAddr,
			}
			conn, err := dialer.Dial("tcp", fmt.Sprintf("%s:%d", ip, m.Port))
			if err == nil {
				conn.Close()
				success = true
			}
		}

		if success {
			return true
		}

		if i < m.Retry {
			time.Sleep(time.Duration(m.Timeout) * time.Millisecond / 2)
		}
	}
	return false
}

func (m *Monitor) updateStatus(ip string, current bool) {
	m.statusLock.Lock()
	defer m.statusLock.Unlock()

	previous := m.status[ip]
	if current != previous {
		m.logStatusChange(ip, previous, current)
		m.status[ip] = current
	}
}

func (m *Monitor) logStatusChange(ip string, from, to bool) {
	statusToStr := func(s bool) string {
		if s {
			return "UP"
		}
		return "DOWN"
	}

	var hostname string
	if m.WithHostname {
		names, err := net.LookupAddr(ip)
		if err == nil && len(names) > 0 {
			hostname = names[0]
		}
	}

	starlog.Infof("[Status Change] %s (%s): %s → %s\n",
		ip,
		hostname,
		statusToStr(from),
		statusToStr(to),
	)
}

func (m *Monitor) displayStatus() {
	fmt.Print("\033[H\033[2J") // Clear screen
	fmt.Printf("Monitoring Status (%s)\n", time.Now().Format("2006-01-02 15:04:05"))
	fmt.Println("======================================")

	m.statusLock.RLock()
	defer m.statusLock.RUnlock()

	for _, ip := range m.IPs {
		status := m.status[ip]
		statusStr := "DOWN"
		if status {
			statusStr = "UP"
		}

		var hostname string
		if m.WithHostname {
			names, err := net.LookupAddr(ip)
			if err == nil && len(names) > 0 {
				hostname = names[0]
			}
		}

		fmt.Printf("%-15s", ip)
		if m.WithHostname {
			fmt.Printf(" (%s)", hostname)
		}
		if statusStr == "UP" {
			starlog.Green(": [%s]\n", statusStr)
		} else {
			starlog.Red(": [%s]\n", statusStr)
		}
	}
}