package main

import (
	"crypto/tls"
	_ "embed"
	"encoding/json"
	"github.com/digineo/go-ping"
	"github.com/gorilla/mux"
	"gopkg.in/yaml.v3"
	"log"
	"net"
	"net/http"
	"net/netip"
	"os"
	"sync"
	"time"
)

type ConfigT struct {
	BindAddress             string        `yaml:"bind_address"`
	ScanSubnets             []string      `yaml:"scan_subnets"`
	DelayBetweenScans       time.Duration `yaml:"delay_between_scans"`
	ExcludeSpecialAddresses bool          `yaml:"exclude_special_addresses"`
	PersistenceLocation     string        `yaml:"persistence_location"`
	UseTLS                  bool          `yaml:"use_tls"`
	TLSKeyFile              string        `yaml:"tls_key_file"`
	TLSCertFile             string        `yaml:"tls_cert_file"`
}

type HostT struct {
	FirstSeen time.Time `json:"first_seen"`
	LastSeen  time.Time `json:"last_seen"`
	Online    bool      `json:"online"`
	RevDNS    string    `json:"rev_dns"`
}

type SubnetT struct {
	Hosts          map[string]HostT `json:"hosts"`
	TotalAddresses int              `json:"total_addresses"`
	UsedAddresses  int              `json:"used_addresses"`
	HostList       []string         `json:"host_list"` //I _REALLY_ didn't want to calculate this in javascript
}

type PersistenceT struct {
	Subnets map[string]SubnetT
}

var conf ConfigT
var p PersistenceT

//go:embed index.html
var indexhtml []byte
var pinger *ping.Pinger

func main() {

	if len(os.Args) < 2 {
		log.Fatalf("Usage: %s <config location>", os.Args[0])
	}

	f, err := os.Open(os.Args[1])
	if err != nil {
		log.Fatalf("Failed to open %s: %s", os.Args[1], err)
	}
	err = yaml.NewDecoder(f).Decode(&conf)
	if err != nil {
		log.Fatalf("Failed to load config file: %s", err)
	}
	_ = f.Close()

	f, err = os.Open(conf.PersistenceLocation)
	if err != nil {
		log.Printf("Failed to open persistence file: %s. Will try to create new one later.", err)
	} else {
		err = json.NewDecoder(f).Decode(&p)
		_ = f.Close()
		if err != nil {
			log.Fatalf("Failed to decode persistence file: %s", err)
		}
	}
	if p.Subnets == nil {
		p.Subnets = make(map[string]SubnetT)
	}

	pinger, err = ping.New("0.0.0.0", "")
	if err != nil {
		log.Fatalf("Failed to create pinger: %s", err)
	}

	go scanner()
	r := mux.NewRouter()
	r.HandleFunc("/", func(writer http.ResponseWriter, request *http.Request) {
		_, _ = writer.Write(indexhtml)
	}).Methods("GET")

	r.HandleFunc("/data", func(writer http.ResponseWriter, request *http.Request) {
		_ = json.NewEncoder(writer).Encode(&p)
	}).Methods("GET")
	server := &http.Server{
		Addr:    conf.BindAddress,
		Handler: r,
		TLSConfig: &tls.Config{
			MinVersion: tls.VersionTLS13,
		},
	}
	if conf.UseTLS {
		err = server.ListenAndServeTLS(conf.TLSCertFile, conf.TLSKeyFile)
	} else {
		err = server.ListenAndServe()
	}
	if err != nil {
		log.Fatalf("%s", err)
	}

}

func PingFunc(addr net.IPAddr) bool {
	_, err := pinger.PingAttempts(&addr, 200*time.Millisecond, 2)
	return err == nil
}

func scanner() {
	for {
		var wg sync.WaitGroup
		var mutex sync.Mutex
		for _, subnet := range conf.ScanSubnets {

			go func(v string) {
				wg.Add(1)
				defer wg.Done()
				persistenceSubnet, ok := p.Subnets[v]
				if !ok {
					persistenceSubnet = SubnetT{}
					persistenceSubnet.Hosts = make(map[string]HostT)
				}
				log.Printf("Scanning subnet %s", v)
				prefix, err := netip.ParsePrefix(v)
				if err != nil {
					log.Printf("Error: %s", err)
					return
				}
				prefix = prefix.Masked()
				addr := prefix.Addr()

				if conf.ExcludeSpecialAddresses {
					addr = addr.Next()
				}
				persistenceSubnet.TotalAddresses = 0
				persistenceSubnet.UsedAddresses = 0
				persistenceSubnet.HostList = make([]string, 0)
				for {
					if !prefix.Contains(addr) {
						break
					}
					//skip broadcast address
					if conf.ExcludeSpecialAddresses && !prefix.Contains(addr.Next()) {
						break
					}
					persistenceSubnet.TotalAddresses++
					//fmt.Println(addr.String())
					persistenceSubnet.HostList = append(persistenceSubnet.HostList, addr.String())
					x := net.IPAddr{IP: net.ParseIP(addr.String())}
					pingstate := PingFunc(x)
					host, ok := persistenceSubnet.Hosts[addr.String()]
					if pingstate {
						persistenceSubnet.UsedAddresses++
						//log.Printf("%s is up", addr.String())
						rdnsString := ""
						rdns, err := net.LookupAddr(addr.String())
						if err == nil {
							if len(rdns) > 0 {
								rdnsString = rdns[0]
							}
						}
						if !ok {
							persistenceSubnet.Hosts[addr.String()] = HostT{
								FirstSeen: time.Now(),
								LastSeen:  time.Now(),
								Online:    true,
								RevDNS:    rdnsString,
							}
						} else {
							host.LastSeen = time.Now()
							host.Online = true
							host.RevDNS = rdnsString
							persistenceSubnet.Hosts[addr.String()] = host
						}

					} else if ok {
						host.Online = false
						persistenceSubnet.Hosts[addr.String()] = host
						persistenceSubnet.UsedAddresses++
					}

					addr = addr.Next()
				}
				mutex.Lock()
				p.Subnets[v] = persistenceSubnet
				mutex.Unlock()
				log.Printf("Scan of %s finished", v)
			}(subnet)
		}
		wg.Wait()
		log.Printf("All scans finished")
		f, err := os.Create(conf.PersistenceLocation)
		if err != nil {
			log.Printf("Failed to save persistence: %s", err)
		} else {
			err := json.NewEncoder(f).Encode(&p)
			_ = f.Close()
			if err != nil {
				log.Printf("Failed to encode persistence: %s", err)
			}
		}
		time.Sleep(conf.DelayBetweenScans)
	}
}