ProxmoxDash/network_flows.go

package main

import (
	"database/sql"
	_ "embed"
	"encoding/json"
	"errors"
	"io"
	"net"
	"net/http"
	"time"
)

//go:embed assets/network.html
var networkHTML string

// Original and reply tuples describe one conntrack entry, not two connections.
type NetworkFlow struct {
	External             bool   `json:"external"`
	Protocol             string `json:"protocol"`
	Source               string `json:"source"`
	Destination          string `json:"destination"`
	SourcePort           int    `json:"sourcePort"`
	DestinationPort      int    `json:"destinationPort"`
	ReplySource          string `json:"replySource"`
	ReplyDestination     string `json:"replyDestination"`
	ReplySourcePort      int    `json:"replySourcePort"`
	ReplyDestinationPort int    `json:"replyDestinationPort"`
	State                string `json:"state"`
	Country              string `json:"country"`
	DestinationCountry   string `json:"destinationCountry"`
	Packets              uint64 `json:"packets"`
	ReplyPackets         uint64 `json:"replyPackets"`
	Bytes                uint64 `json:"bytes"`
	ReplyBytes           uint64 `json:"replyBytes"`
	Accounting           bool   `json:"accounting"`
	Zone                 string `json:"zone,omitempty"`
}
type FirewallCounter struct {
	Name    string `json:"name"`
	Packets uint64 `json:"packets"`
	Bytes   uint64 `json:"bytes"`
}
type NetworkSnapshot struct {
	At            int64             `json:"at"`
	ReceivedAt    int64             `json:"receivedAt"`
	Flows         []NetworkFlow     `json:"flows"`
	Error         string            `json:"error,omitempty"`
	Limited       bool              `json:"limited"`
	ParseErrors   int               `json:"parseErrors"`
	Firewall      []FirewallCounter `json:"firewall"`
	FirewallAt    int64             `json:"firewallAt"`
	FirewallError string            `json:"firewallError,omitempty"`
	AgentName     string            `json:"agentName"`
}

func networkExternalIP(value string) bool {
	ip := net.ParseIP(value)
	if ip == nil || !ip.IsGlobalUnicast() || ip.IsPrivate() || ip.IsLoopback() || ip.IsLinkLocalUnicast() {
		return false
	}
	if v4 := ip.To4(); v4 != nil && v4[0] == 100 && v4[1] >= 64 && v4[1] <= 127 {
		return false
	}
	return true
}
func (s *Store) SaveNetwork(agent string, snapshot NetworkSnapshot) error {
	if len(snapshot.Flows) > 5000 || len(snapshot.Firewall) > 128 {
		return errors.New("network snapshot too large")
	}
	if snapshot.At < time.Now().Unix()-120 || snapshot.At > time.Now().Unix()+60 {
		return errors.New("invalid snapshot time")
	}
	for i := range snapshot.Flows {
		f := &snapshot.Flows[i]
		if net.ParseIP(f.Source) == nil || net.ParseIP(f.Destination) == nil || net.ParseIP(f.ReplySource) == nil || net.ParseIP(f.ReplyDestination) == nil {
			return errors.New("invalid network address")
		}
		for _, p := range []int{f.SourcePort, f.DestinationPort, f.ReplySourcePort, f.ReplyDestinationPort} {
			if p < 0 || p > 65535 {
				return errors.New("invalid network port")
			}
		}
		if f.Protocol != "tcp" && f.Protocol != "udp" && f.Protocol != "icmp" && f.Protocol != "icmpv6" {
			return errors.New("invalid protocol")
		}
		f.State = trafficText(f.State, 32)
		f.Zone = trafficText(f.Zone, 16)
		f.External = networkExternalIP(f.Source) || networkExternalIP(f.Destination)
		f.Country = lookupTrafficCountry(f.Source)
		f.DestinationCountry = lookupTrafficCountry(f.Destination)
	}
	snapshot.Error = trafficText(snapshot.Error, 600)
	snapshot.FirewallError = trafficText(snapshot.FirewallError, 600)
	for i := range snapshot.Firewall {
		snapshot.Firewall[i].Name = trafficText(snapshot.Firewall[i].Name, 200)
	}
	snapshot.ReceivedAt = time.Now().Unix()
	snapshot.AgentName = ""
	raw, err := json.Marshal(snapshot)
	if err != nil {
		return err
	}
	_, err = s.db.Exec(`INSERT INTO network_sources(agent,snapshot) VALUES(?,?) ON CONFLICT(agent) DO UPDATE SET snapshot=excluded.snapshot`, agent, string(raw))
	return err
}
func networkRoutes(mux *http.ServeMux, s *Store) {
	enabled := func() bool { c := s.TrafficSettings(); return c.Enabled && c.NetworkEnabled }
	mux.HandleFunc("/network", func(w http.ResponseWriter, r *http.Request) {
		if !enabled() {
			http.Error(w, "Включите TCP/UDP в настройках карты трафика", 404)
			return
		}
		w.Header().Set("Content-Type", "text/html; charset=utf-8")
		w.Write([]byte(networkHTML))
	})
	mux.HandleFunc("/api/network", func(w http.ResponseWriter, r *http.Request) {
		if r.Method != "GET" {
			w.WriteHeader(405)
			return
		}
		if !enabled() {
			http.Error(w, "Сбор сети отключён", 404)
			return
		}
		var raw, name string
		err := s.db.QueryRow(`SELECT n.snapshot,a.name FROM network_sources n JOIN agents a ON a.id=n.agent WHERE n.agent=?`, s.TrafficSettings().AgentID).Scan(&raw, &name)
		if err != nil && err != sql.ErrNoRows {
			writeJSON(w, nil, err)
			return
		}
		if err != nil {
			writeJSON(w, NetworkSnapshot{Flows: []NetworkFlow{}, Error: "Ожидание агента v0.21.0 или новее"}, nil)
			return
		}
		var snapshot NetworkSnapshot
		if err = json.Unmarshal([]byte(raw), &snapshot); err != nil {
			writeJSON(w, nil, err)
			return
		}
		snapshot.AgentName = name
		writeJSON(w, snapshot, nil)
	})
	mux.HandleFunc("/api/agent/network", func(w http.ResponseWriter, r *http.Request) {
		if r.Method != "POST" {
			w.WriteHeader(405)
			return
		}
		id, secret, ok := agentCredentials(r)
		if !ok || !s.AuthenticateAgent(id, secret) {
			http.Error(w, "Unauthorized", 401)
			return
		}
		c := s.TrafficSettings()
		if !c.Enabled || !c.NetworkEnabled || c.AgentID != id {
			http.Error(w, "Источник отключён", 403)
			return
		}
		var snapshot NetworkSnapshot
		if json.NewDecoder(io.LimitReader(r.Body, 4<<20)).Decode(&snapshot) != nil {
			http.Error(w, "Invalid snapshot", 400)
			return
		}
		if err := s.SaveNetwork(id, snapshot); err != nil {
			http.Error(w, err.Error(), 400)
			return
		}
		writeJSON(w, map[string]bool{"accepted": true}, nil)
	})
}