package main

import (
	"bufio"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"path"
	"path/filepath"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"

	"codeberg.org/eduVPN/vpn-daemon/wgctrl"
	"codeberg.org/eduVPN/vpn-daemon/wgctrl/wgtypes"
)

var (
	credDirEnvVar       = "CREDENTIALS_DIRECTORY"
	socketDir           = "/run/openvpn-server"
	maintenanceModePath = "/run/vpn-daemon/maintenance-mode"
)

type NodeInfo struct {
	RelLoadAverage  []int     `json:"rel_load_average"`
	LoadAverage     []float64 `json:"load_average"`
	CPUCount        int       `json:"cpu_count"`
	NodeUptime      int       `json:"node_uptime"`
	MaintenanceMode bool      `json:"maintenance_mode"`
	OClientCount    *int      `json:"o_client_count,omitempty"`
	WPeerCount      *int      `json:"w_peer_count,omitempty"`
}

type VpnClientInfo struct {
	CommonName string `json:"common_name"`
	IPFour     string `json:"ip_four"`
	IPSix      string `json:"ip_six"`
}

type WgPeerInfo struct {
	PublicKey         string        `json:"public_key"`
	IPNet             VpnNetList    `json:"ip_net"`
	LastHandshakeTime HandshakeTime `json:"last_handshake_time"`
	BytesIn           int64         `json:"bytes_in"`
	BytesOut          int64         `json:"bytes_out"`
}

type (
	HandshakeTime time.Time
	VpnNetList    []net.IPNet
)

// we want to marshal "IsZero" time to JSON null
func (t HandshakeTime) MarshalJSON() ([]byte, error) {
	if !time.Time(t).IsZero() {
		return json.Marshal(time.Time(t).UTC().Format(time.RFC3339))
	}
	return json.Marshal(nil)
}

// marshal to a string list in IP/prefix notation
func (vpnNetList VpnNetList) MarshalJSON() ([]byte, error) {
	netList := make([]string, 0, len(vpnNetList))
	for _, ipNet := range vpnNetList {
		netList = append(netList, ipNet.String())
	}
	return json.Marshal(netList)
}

func oConnectionList() []VpnClientInfo {
	// ls /run/openvpn-server/*.sock | wc
	socketList, err := filepath.Glob(socketDir + "/*.sock")
	if err != nil {
		log.Print(err)
		return []VpnClientInfo{}
	}

	vpnClientInfoChannel := make(chan []*VpnClientInfo, len(socketList))
	for socketIndex := 0; socketIndex < len(socketList); socketIndex++ {
		go obtainStatus(socketList[socketIndex], vpnClientInfoChannel)
	}

	connectionList := make([]VpnClientInfo, 0)
	for socketIndex := 0; socketIndex < len(socketList); socketIndex++ {
		vpnClientInfoList := <-vpnClientInfoChannel
		for _, vpnClientInfo := range vpnClientInfoList {
			connectionList = append(connectionList, *vpnClientInfo)
		}
	}

	return connectionList
}

func ovList() http.HandlerFunc {
	return func(w http.ResponseWriter, _ *http.Request) {
		sendJSONResponse(w, http.StatusOK, map[string][]VpnClientInfo{"connection_list": oConnectionList()})
	}
}

func validateConnectionID(connectionID string) (*string, error) {
	// @see https://lore.kernel.org/wireguard/X+UkseUOEY1sVDEe@zx2c4.com/
	connectionIDRegexp := regexp.MustCompile(`^[A-Za-z0-9+\\/]{42}[A|E|I|M|Q|U|Y|c|g|k|o|s|w|4|8|0]=$`)
	if !connectionIDRegexp.MatchString(connectionID) {
		return nil, errors.New("invalid connection_id")
	}

	return &connectionID, nil
}

func ovDisconnect() http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		commonName, err := validateConnectionID(r.FormValue("common_name"))
		if err != nil {
			sendJSONErrorResponse(w, http.StatusBadRequest, "invalid \"common_name\"")
			return
		}

		// ls /run/openvpn-server/*.sock | wc
		socketList, err := filepath.Glob(socketDir + "/*.sock")
		if err != nil {
			sendJSONErrorResponse(w, http.StatusInternalServerError, "unable to list file system sockets: "+err.Error())
			return
		}

		// as "DISCONNECT" has no need to pass information back here, we
		// use a WaitGroup instead of a channel...
		var wg sync.WaitGroup
		for socketIndex := 0; socketIndex < len(socketList); socketIndex++ {
			wg.Add(1)
			go disconnectClient(socketList[socketIndex], *commonName, &wg)
		}
		wg.Wait()

		// HTTP 204 (No Content)
		w.WriteHeader(http.StatusNoContent)
	}
}

func wgAddPeer(deviceName string) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		publicKey, err := validateConnectionID(r.FormValue("public_key"))
		if err != nil {
			sendJSONErrorResponse(w, http.StatusBadRequest, "invalid \"public_key\"")
			return
		}

		pK, err := wgtypes.ParseKey(*publicKey)
		if err != nil {
			sendJSONErrorResponse(w, http.StatusBadRequest, "unable to parse \"public_key\"")
			return
		}

		ipNets := r.PostForm["ip_net"]
		netList := make([]net.IPNet, 0, len(ipNets))
		for _, ipNetString := range ipNets {
			_, ipNet, err := net.ParseCIDR(ipNetString)
			if err != nil {
				sendJSONErrorResponse(w, http.StatusBadRequest, "invalid \"ip_net\"")
				return
			}
			netList = append(netList, *ipNet)
		}
		if len(netList) == 0 {
			sendJSONErrorResponse(w, http.StatusBadRequest, "\"ip_net\" did not contain any valid IP addresses")
			return
		}

		c, err := wgctrl.New()
		if err != nil {
			log.Print(err)
			sendJSONErrorResponse(w, http.StatusInternalServerError, "wgctrl.New() failed")
			return
		}
		defer c.Close()

		wgConfig := wgtypes.Config{
			ReplacePeers: false,
			Peers: []wgtypes.PeerConfig{{
				PublicKey:         pK,
				Remove:            false,
				UpdateOnly:        false,
				ReplaceAllowedIPs: true,
				AllowedIPs:        netList,
			}},
		}

		if c.ConfigureDevice(deviceName, wgConfig) != nil {
			log.Printf("(*Client).ConfigureDevice(%s) failed: %s", deviceName, err)
			sendJSONErrorResponse(w, http.StatusInternalServerError, fmt.Sprintf("(*Client).ConfigureDevice(%s) failed", deviceName))
			return
		}

		// HTTP 204 (No Content)
		w.WriteHeader(http.StatusNoContent)
	}
}

func wgRemovePeer(deviceName string) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		publicKey, err := validateConnectionID(r.FormValue("public_key"))
		if err != nil {
			sendJSONErrorResponse(w, http.StatusBadRequest, "invalid \"public_key\"")
			return
		}

		c, err := wgctrl.New()
		if err != nil {
			log.Print(err)
			sendJSONErrorResponse(w, http.StatusInternalServerError, "wgctrl.New() failed")
			return
		}
		defer c.Close()
		d, err := c.Device(deviceName)
		if err != nil {
			log.Printf("(*Client).Device(%s) failed: %s", deviceName, err)
			sendJSONErrorResponse(w, http.StatusInternalServerError, fmt.Sprintf("(*Client).Device(%s) failed", deviceName))
			return
		}

		pK, err := wgtypes.ParseKey(*publicKey)
		if err != nil {
			sendJSONErrorResponse(w, http.StatusBadRequest, "unable to parse \"public_key\"")
			return
		}

		// find peer
		for _, p := range d.Peers {
			if *publicKey == p.PublicKey.String() {
				// found!
				wgConfig := wgtypes.Config{
					ReplacePeers: false,
					Peers: []wgtypes.PeerConfig{{
						PublicKey:         pK,
						Remove:            true,
						ReplaceAllowedIPs: true,
					}},
				}

				// delete peer
				c.ConfigureDevice(deviceName, wgConfig)

				// create info to return in JSON response
				peerInfo := WgPeerInfo{
					PublicKey:         p.PublicKey.String(),
					IPNet:             VpnNetList(p.AllowedIPs),
					LastHandshakeTime: HandshakeTime(p.LastHandshakeTime),
					BytesIn:           p.ReceiveBytes, BytesOut: p.TransmitBytes,
				}

				sendJSONResponse(w, http.StatusOK, peerInfo)
				return
			}
		}

		// We were unable to find peer, but that is okay!
		// HTTP 204 (No Content)
		w.WriteHeader(http.StatusNoContent)
	}
}

func wPeerList(deviceName string, showAll bool) []WgPeerInfo {
	c, err := wgctrl.New()
	if err != nil {
		log.Print(err)
		return []WgPeerInfo{}
	}
	defer c.Close()

	peerList := make([]WgPeerInfo, 0)

	d, err := c.Device(deviceName)
	if err != nil {
		log.Print(err)
		return []WgPeerInfo{}
	}

	for _, p := range d.Peers {
		if !showAll {
			// ignore peers never seen
			if p.LastHandshakeTime.IsZero() {
				continue
			}
			// ignore peers seen > 3 minutes ago
			if !p.LastHandshakeTime.Add(time.Minute * 3).After(time.Now()) {
				continue
			}
		}

		peerInfo := WgPeerInfo{
			PublicKey:         p.PublicKey.String(),
			IPNet:             VpnNetList(p.AllowedIPs),
			LastHandshakeTime: HandshakeTime(p.LastHandshakeTime),
			BytesIn:           p.ReceiveBytes,
			BytesOut:          p.TransmitBytes,
		}

		peerList = append(peerList, peerInfo)
	}

	return peerList
}

func wgList(deviceName string) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		showAll := r.FormValue("show_all") == "yes"
		sendJSONResponse(w, http.StatusOK, map[string][]WgPeerInfo{"peer_list": wPeerList(deviceName, showAll)})
	}
}

func obtainStatus(socketFile string, vpnClientInfoChannel chan []*VpnClientInfo) {
	managementConnection, err := getManagementConnection(socketFile)
	if err != nil {
		vpnClientInfoChannel <- []*VpnClientInfo{}
		return
	}
	defer managementConnection.Close()

	vpnClientInfoList := []*VpnClientInfo{}
	fmt.Fprintf(managementConnection, "status 2\n")
	managementPortScanner := bufio.NewScanner(managementConnection)
	for managementPortScanner.Scan() {
		text := managementPortScanner.Text()
		if strings.Index(text, "END") == 0 {
			break
		}
		if strings.Index(text, "CLIENT_LIST") == 0 {
			// HEADER,CLIENT_LIST,Common Name,Real Address,Virtual Address,
			//      Virtual IPv6 Address,Bytes Received,Bytes Sent,
			//      Connected Since,Connected Since (time_t),Username,
			//      Client ID,Peer ID
			strList := strings.Split(text, ",")
			if strList[1] != "UNDEF" && strList[3] != "" && strList[4] != "" {
				// only add clients with CN != "UNDEF" and IP addresses are not
				// empty strings...
				vpnClientInfoList = append(vpnClientInfoList, &VpnClientInfo{strList[1], strList[3], strList[4]})
			}
		}
	}

	vpnClientInfoChannel <- vpnClientInfoList
}

func disconnectClient(socketFile string, commonName string, wg *sync.WaitGroup) {
	defer wg.Done()
	managementConnection, err := getManagementConnection(socketFile)
	if err != nil {
		return
	}
	defer managementConnection.Close()

	managementPortScanner := bufio.NewScanner(managementConnection)
	fmt.Fprintf(managementConnection, "kill %s\n", commonName)
	for managementPortScanner.Scan() {
		text := managementPortScanner.Text()
		if strings.Index(text, "ERROR") == 0 || strings.Index(text, "SUCCESS") == 0 {
			// we do not care about the actual result, we can't do anything
			// anyway...
			break
		}
	}
}

func getManagementConnection(socketFile string) (net.Conn, error) {
	managementConnection, err := net.Dial("unix", socketFile)
	if err != nil {
		log.Printf("WARNING: %s", err)
		return nil, err
	}

	return managementConnection, nil
}

func sendJSONErrorResponse(w http.ResponseWriter, responseCode int, errorMessage string) error {
	return sendJSONResponse(w, responseCode, map[string]string{"error": errorMessage})
}

func sendJSONResponse(w http.ResponseWriter, responseCode int, responseData interface{}) error {
	r, err := json.Marshal(responseData)
	if err != nil {
		return err
	}
	w.Header().Add("Content-Type", "application/json")
	w.WriteHeader(responseCode)
	w.Write(r)
	return nil
}

func nodeInfo(wgDevice string) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		includeClientPeerCount := r.FormValue("include_client_peer_count") == "yes"

		// $ cat /proc/loadavg
		// 0.95 0.82 0.68 2/996 93945
		loadAvgList := []float64{}
		relLoadAvgList := []int{}
		procLoadAvg, err := os.ReadFile("/proc/loadavg")
		if err != nil {
			// we were unable to open the file, perhaps because of permissions
			// or because this is not Linux. In either case, pretend we got an
			// empty file
			procLoadAvg = []byte{}
		}
		procLoadAvgFields := strings.Fields(string(procLoadAvg))
		if len(procLoadAvgFields) >= 3 {
			for _, a := range procLoadAvgFields[0:3] {
				loadAvg, err := strconv.ParseFloat(a, 64)
				if err != nil {
					// if we can't parse it as a float, then just ignore it
					continue
				}
				loadAvgList = append(loadAvgList, loadAvg)
				relLoadAvgList = append(relLoadAvgList, int(loadAvg*100)/runtime.NumCPU())
			}
		}

		// $ cat /proc/uptime
		// 3182.25 5256.60
		nodeUptime := 0
		procUptime, err := os.ReadFile("/proc/uptime")
		if err != nil {
			// we were unable to open the file, perhaps because of permissions
			// or because this is not Linux. In either case, pretend we got an
			// empty file
			procUptime = []byte{}
		}
		procUptimeFields := strings.Fields(string(procUptime))
		if len(procUptimeFields) >= 2 {
			sysUptime, err := strconv.ParseFloat(procUptimeFields[0], 64)
			if err == nil {
				nodeUptime = int(sysUptime)
			}
		}

		// Check if path exists
		maintenanceMode := false
		_, err = os.Stat(maintenanceModePath)
		if err == nil {
			maintenanceMode = true
		} else if !os.IsNotExist(err) {
			log.Print(err)
		}

		var oClientCount *int
		var wPeerCount *int

		if includeClientPeerCount {
			oConnectionListLen := len(oConnectionList())
			wPeerListLen := len(wPeerList(wgDevice, false))
			oClientCount = &oConnectionListLen
			wPeerCount = &wPeerListLen
		}

		nodeInfo := NodeInfo{
			RelLoadAverage:  relLoadAvgList,
			LoadAverage:     loadAvgList,
			CPUCount:        runtime.NumCPU(),
			NodeUptime:      nodeUptime,
			MaintenanceMode: maintenanceMode,
			OClientCount:    oClientCount,
			WPeerCount:      wPeerCount,
		}

		sendJSONResponse(w, http.StatusOK, nodeInfo)
	}
}

func getTLSConfig(clientCaCertFile string) *tls.Config {
	clientCaCertPem, err := os.ReadFile(clientCaCertFile)
	fatalIfError(err)
	clientCaCertPool := x509.NewCertPool()
	if !clientCaCertPool.AppendCertsFromPEM(clientCaCertPem) {
		log.Fatal(errors.New("unable to register client CA certificate"))
	}

	return &tls.Config{
		MinVersion: tls.VersionTLS13,
		ClientAuth: tls.RequireAndVerifyClientCert,
		ClientCAs:  clientCaCertPool,
	}
}

func main() {
	hostPort := flag.String("listen", "127.0.0.1:41194", "IP:port to listen on")
	wgDevice := flag.String("wg-device", "wg0", "WireGuard device")

	flag.Usage = func() {
		flag.PrintDefaults()
	}
	flag.Parse()

	http.HandleFunc("GET /i/node", nodeInfo(*wgDevice))
	http.HandleFunc("POST /w/add_peer", wgAddPeer(*wgDevice))
	http.HandleFunc("POST /w/remove_peer", wgRemovePeer(*wgDevice))
	http.HandleFunc("GET /w/peer_list", wgList(*wgDevice))
	http.HandleFunc("GET /o/connection_list", ovList())
	http.HandleFunc("POST /o/disconnect_client", ovDisconnect())

	credDir := os.Getenv(credDirEnvVar)
	if credDir != "" {
		log.Printf("'%s' set to '%s', enabling TLS", credDirEnvVar, credDir)

		// if a credentials directory is specified, enable TLS
		clientCaCertFile := path.Join(credDir, "ca.crt")
		serverCertFile := path.Join(credDir, "server.crt")
		serverKeyFile := path.Join(credDir, "server.key")

		server := &http.Server{
			Addr:      *hostPort,
			TLSConfig: getTLSConfig(clientCaCertFile),
		}
		log.Fatal(server.ListenAndServeTLS(serverCertFile, serverKeyFile))

		return
	}

	log.Fatal(http.ListenAndServe(*hostPort, nil))
}

func fatalIfError(err error) {
	if err != nil {
		log.Fatal(err)
	}
}
