hvpn-node3/cmd/hvpn-node/hvpn-node.go

package main

import (
	"bufio"
	"fmt"
	"log/slog"
	"net/http"
	"net/netip"
	"os"
	"os/signal"
	"strings"
	"time"

	"github.com/urfave/cli/v2"
	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"

	hvpnnode3 "gitea.hbanafa.com/HeshamTB/hvpn-node3"
)

/*
   Can change all therse vars to local func vars
   IPPool can be internal to wglink or other abstraction. As well as
   Port, ifname and CIDR and so on. Use Clie.ctx to get the values
*/

var IPPool hvpnnode3.IPPool
var VPNIPCIDR string
var PrivateKeyPath cli.Path
var InterfaceName string
var WgPort int
var wgLink *hvpnnode3.WGLink

var httpPort int

func main() {


    // TODO: Define error exit codes
    slog.SetLogLoggerLevel(slog.LevelDebug)


    app := createCliApp()
    err := app.Run(os.Args)
    if err != nil {
        slog.Error(err.Error())
        os.Exit(1)
    }
}

func run(ctx *cli.Context) {
    slog.Debug("Starting run()")

    apiMux := http.NewServeMux()
    apiMux.HandleFunc("GET /node", hvpnnode3.HandleGetNodeInfo(wgLink))
    apiMux.HandleFunc("GET /peer/{pubkey}", hvpnnode3.HandleGetPeer(wgLink))
    apiMux.HandleFunc("POST /peer", hvpnnode3.HandlePostPeer(wgLink))
    apiMux.HandleFunc("DELETE /peer/{pubkey}", hvpnnode3.HandleDeletePeer(wgLink))
    apiMux.HandleFunc("GET /peers", hvpnnode3.HandleGetPeers(wgLink))

    var handler http.Handler = apiMux
    handler = hvpnnode3.HttpLogHandler2(handler)

    port := fmt.Sprintf("%d", httpPort)
    host := ctx.String("host")
    if host == "" {
        slog.Info("Host is not set. Using 0.0.0.0")
        host = "0.0.0.0"
    }
    hostPort := fmt.Sprintf("%s:%s", host, port)

    slog.Info(fmt.Sprintf("Starting HVPN node on %s", hostPort))
    srv := &http.Server{
        ReadTimeout: 120 * time.Second,
        WriteTimeout: 120 * time.Second,
        IdleTimeout: 120 * time.Second,
        Handler: handler,
        Addr: hostPort,
    }

    defer wgLink.Close()
    slog.Warn(srv.ListenAndServe().Error())
}

func createCliApp() *cli.App {
    app := cli.NewApp()
    app.Name = os.Args[0]
    app.HideHelpCommand = true
    app.Usage = "HVPN node API server"
    app.Args = true

    author1 := cli.Author{
        Name: "Hesham T. Banafa",
        Email: "hishaminv@gmail.com"}
    app.Authors = append(app.Authors, &author1)
    
    logLevel := cli.StringFlag{
        Name: "log-level",
        Value: "INFO",
        EnvVars: []string{"LOG_LEVEL"},
        Action: func(ctx *cli.Context, s string) error {
            lvl := new(slog.LevelVar)
            err := lvl.UnmarshalText([]byte(s))
            if err != nil {
                slog.Debug("Error on unmarshal log level flag")
                return err
            }
            slog.SetLogLoggerLevel(lvl.Level())
            return nil 
        },
    }
    app.Flags = append(app.Flags, &logLevel)

    privateKeyFileFlag := cli.PathFlag{
        Name: "private-key",
        Usage: "Path to file with private key",
        Destination: &PrivateKeyPath,
    }
    app.Flags = append(app.Flags, &privateKeyFileFlag)

    vpnIpCIDR := cli.StringFlag{
        Name: "cidr",
        Usage: "The network subnet used for the internal IP Pool",
        Value: "10.42.0.0/16",
        Aliases: []string{"n"},
        Destination: &VPNIPCIDR,
    }
    app.Flags = append(app.Flags, &vpnIpCIDR)

    wgInterfaceName := cli.StringFlag{
        Name: "interface",
        Usage: "Name of the Wireguard interface to be created and managed",
        Value: "hvpn0",
        Aliases: []string{"i"},
        Destination: &InterfaceName,
    }
    app.Flags = append(app.Flags, &wgInterfaceName)

    wgEndpoint := cli.StringFlag{
        Name: "endpoint",
        Usage: "Wireguard endpoint domain or address without the port",
        Value: "domain.name.notset",
    }
    app.Flags = append(app.Flags, &wgEndpoint)

    wgPort := cli.IntFlag{
        Name: "port",
        Usage: "UDP Port for wireguard device",
        Value: 6416,
        Aliases: []string{"p"},
        Destination: &WgPort,
    }
    app.Flags = append(app.Flags, &wgPort)

    httpListenAddr := cli.StringFlag{
        Name: "host",
        Usage: "IP address to listen on for HTTP API requests",
        Value: "0.0.0.0",
        Action: func(ctx *cli.Context, s string) error {
            _, err := netip.ParseAddr(s)
            if err != nil {
                return cli.Exit(fmt.Sprintf("Can not parse %s as a network IP", s), 1)
            }
            return nil
        },
    }
    app.Flags = append(app.Flags, &httpListenAddr)

    httpPort := cli.IntFlag{
        Name: "http-port",
        Usage: "TCP Port for HTTP API",
        Value: 8080,
        Destination: &httpPort,
    }
    app.Flags = append(app.Flags, &httpPort)


    app.Action = func(ctx *cli.Context) error {
        err := setup(ctx)
        if err != nil {
            return err
        }
        run(ctx)
        return nil
    }

    return app
}

func setup(ctx *cli.Context) error {
    slog.Debug("Starting setup()")
        uid := os.Getuid()
    if uid == -1 {
        slog.Warn("Running on windows! whatrr u doing?")
    } else if uid == 0 {
        slog.Warn("Running as root! avoid running as root by setting CAP_NET_ADMIN")
    }
    uuid, err := hvpnnode3.InitNodeUUID()
    if err != nil {
        slog.Error(err.Error())
        os.Exit(-1)
    }
    slog.Info("Node UUID: " + uuid.String())

    var privateKey wgtypes.Key
    createPrivKey := func() error {
        slog.Info("Creating a private key")
        privateKey, err = wgtypes.GeneratePrivateKey()
        if err != nil {
            return cli.Exit(err, 1)
        }
        slog.Debug(fmt.Sprintf("new public key: %s", privateKey.PublicKey().String()))
        return nil
    }

    if PrivateKeyPath == "" {
        err := createPrivKey()
        if err != nil {
            return err
        }
    } else {
        privKeyFile, err := os.Open(PrivateKeyPath)
        defer privKeyFile.Close()
        if err != nil {
            slog.Error(err.Error())
            slog.Info("Could not open private key file")
            err := createPrivKey()
            if err != nil {
                return err
            }
        } else {
            privateKeyStr := make([]byte, 45)
            n, err := privKeyFile.Read(privateKeyStr)
            if err != nil {
                return cli.Exit(err, 1)
            }
            if n != 45 {
                slog.Warn("Private key length did not math the expected 45!")
            }
            slog.Debug(fmt.Sprintf("Read %d bytes from keyfile", n))

            privateKey, err = wgtypes.ParseKey(string(privateKeyStr))
            slog.Debug("Keyfile opened for reading")
            if err != nil {
                return cli.Exit(err, 1)
            }
            slog.Debug("Private key parsed and is correct")
        }
    }

    wg, err := hvpnnode3.InitWGLink(
        InterfaceName, 
        &privateKey,
        WgPort,
        ctx.String("endpoint"),
    )
    if err != nil {
        slog.Warn("Error while initlizing Wireguard netlink and device!")
        slog.Warn("Ensure to run as root or with CAP_NET_ADMIN")
        return cli.Exit(err, 1)
    }

    wgLink = wg

    slog.Info("Wireguard netlink is setup and running")


    dev, err := wgLink.Device(InterfaceName)
    if err != nil {
        return cli.Exit(err, 1)
    }
    slog.Info("Initialized Wireguard device using wgctrl")

    slog.Info(
        fmt.Sprintf(
            "Device name: %s, UDP port: %d, Type: %s",
            dev.Name, dev.ListenPort, dev.Type.String(),
        ),
    )

    ipPool, err := hvpnnode3.NewPool(VPNIPCIDR)
    if err != nil {
        slog.Error(fmt.Sprintf("IPPool: %s", err))
        os.Exit(1)
    }
    slog.Debug(fmt.Sprintf("Init ip pool %s", VPNIPCIDR))

    testVip, err := ipPool.Allocate()
    if err != nil {
        slog.Error("main.testVip: ", err)
        os.Exit(1)
    }

    slog.Debug(fmt.Sprintf("IP Pool Test IP: %s", testVip.String()))
    err = ipPool.Free(testVip)
    if err != nil {
        slog.Error("Could not free test Vip from IPPool!", err)
        return cli.Exit(err.Error(), 1)
    }
    slog.Debug("Test IP Freed")

    IPPool = ipPool
    wgLink.IPPool = ipPool

    //defer wgLink.Close()
    cInput := make(chan struct{})
    go handleStdin(cInput)
    c := make(chan os.Signal, 1)
    signal.Notify(c, os.Interrupt, os.Kill)
    go func() {
        slog.Info("Listening for SIGINT, SIGKILL and user input")
        select {
        case <- c:
            slog.Warn("Recived SIGINT! Closing Wireguard Device")
        case <- cInput:
            slog.Warn("Recieved quit! Closing Wireguard Device")
        }
        wgLink.Close()
        os.Exit(0)
    }()

    err = testWgPeerAdd(wgLink)
    if err != nil {
        return err
    }

    return nil
}

func testWgPeerAdd(wgLink *hvpnnode3.WGLink) error {
    privateKey, err := wgtypes.GeneratePrivateKey()
    if err != nil {
        return err
    }
    publicKey := privateKey.PublicKey()
    _, err = wgLink.AddPeer(publicKey.String())
    if err != nil {
        slog.Error(err.Error())
        return err
    }
    slog.Debug(fmt.Sprintf("Added test peer %v", publicKey.String()))
    
    err = wgLink.DeletePeer(publicKey.String())
    if err != nil {
        return err
    }
    peers, err := wgLink.GetAllPeers()
    if len(peers) != 0 {
        slog.Warn(fmt.Sprintf("Expected 0 peers, got %d", len(peers)))
    }
    slog.Debug("Removed test peer")

    return nil
}

func handleStdin(c chan struct{}) {
    for {
        reader := bufio.NewReader(os.Stdin)
        slog.Info("Enter 'q' or 'exit' to quit")
        in, _ := reader.ReadString('\n')
        in = strings.ReplaceAll(in, "\n", "")
        if in == "q" || in == "exit" {
            c <- struct{}{}
            break
        }
    }
}