awg-profiler-golang/util.go

package main

import (
	"bufio"
	"context"
	crand "crypto/rand"
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"time"
)

func itoa(n int) string     { return strconv.Itoa(n) }
func itoa64(n int64) string { return strconv.FormatInt(n, 10) }

// atoiOrDie parses a base-10 integer, aborting on malformed input.
func atoiOrDie(s string) int {
	n, err := strconv.Atoi(strings.TrimSpace(s))
	if err != nil {
		die("expected integer, got %q", s)
	}
	return n
}

// ─── logging helpers (mirror the bash die/info/warn) ────────────────────────────

// webMode makes die() panic instead of exiting the process, so the long-running
// web server can recover from an operation failure and turn it into an HTTP
// error response rather than crashing. It is set only by the `web` command.
var webMode bool

// dieError carries a die() message across a recover() in web mode.
type dieError struct{ msg string }

func (e dieError) Error() string { return e.msg }

// die prints an error to stderr and terminates with status 1. In web mode it
// panics with a dieError instead, to be recovered by the HTTP handler wrapper.
func die(format string, args ...any) {
	if webMode {
		panic(dieError{fmt.Sprintf(format, args...)})
	}
	fmt.Fprintf(os.Stderr, "ERROR: "+format+"\n", args...)
	os.Exit(1)
}

func info(format string, args ...any) {
	fmt.Printf("INFO: "+format+"\n", args...)
}

func warn(format string, args ...any) {
	fmt.Fprintf(os.Stderr, "WARN: "+format+"\n", args...)
}

// requireBinary aborts unless the named executable is on PATH.
func requireBinary(name string) {
	if _, err := exec.LookPath(name); err != nil {
		die("%s not installed", name)
	}
}

// ─── interactive input ──────────────────────────────────────────────────────────

var stdinReader = bufio.NewReader(os.Stdin)

// confirm asks a yes/no question; returns true only for a bare y/Y.
func confirm(prompt string) bool {
	fmt.Printf("%s [y/N] ", prompt)
	answer, _ := stdinReader.ReadString('\n')
	answer = strings.TrimSpace(answer)
	return answer == "y" || answer == "Y"
}

// promptDefault reads a value, falling back to def when the user hits enter.
func promptDefault(prompt, def string) string {
	fmt.Printf("%s [%s]: ", prompt, def)
	value, _ := stdinReader.ReadString('\n')
	value = strings.TrimRight(value, "\r\n")
	if value == "" {
		return def
	}
	return value
}

// ─── randomness (crypto/rand, replaces od < /dev/urandom) ───────────────────────

func randUint32() uint32 {
	var b [4]byte
	if _, err := crand.Read(b[:]); err != nil {
		die("failed to read random bytes: %v", err)
	}
	return binary.BigEndian.Uint32(b[:])
}

// randMagic returns a random unsigned 32-bit integer in [5, 2^31-1].
func randMagic() int64 {
	const span = 2147483647 - 5 + 1 // inclusive [5, 2^31-1]
	n := uint64(randUint32())
	return int64(n%span) + 5
}

// randRange returns a uniform integer in the inclusive range [min, max].
// If max < min it warns and returns min so the range stays sane.
func randRange(min, max int) int {
	if max < min {
		warn("Randomisation range [%d,%d] invalid (MTU too small?) — using %d", min, max, min)
		return min
	}
	span := uint64(max - min + 1)
	n := uint64(randUint32())
	return int(n%span) + min
}

func minInt(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func maxInt(a, b int) int {
	if a > b {
		return a
	}
	return b
}

// ─── name sanitisation ──────────────────────────────────────────────────────────

var sanitizeRe = regexp.MustCompile(`[^a-zA-Z0-9._-]`)

func sanitize(v string) string {
	return sanitizeRe.ReplaceAllString(v, "_")
}

// ─── IP helpers ─────────────────────────────────────────────────────────────────

// networkBase strips the /CIDR suffix, returning the bare address portion.
func networkBase(network string) string {
	if i := strings.Index(network, "/"); i >= 0 {
		return network[:i]
	}
	return network
}

// requireSlash24 validates that network is an IPv4 CIDR with a /24 prefix.
// serverIP/getFirstClientIP/incrementIP only ever vary the last octet, and
// every written config hardcodes "Address = <ip>/24" regardless of what was
// typed in — so any other prefix would silently produce a broken, internally
// inconsistent config rather than the network the operator actually asked for.
func requireSlash24(network string) {
	ip, ipnet, err := net.ParseCIDR(network)
	if err != nil || ip.To4() == nil {
		die("Invalid network %q: expected an IPv4 CIDR, e.g. 10.0.0.0/24", network)
	}
	ones, bits := ipnet.Mask.Size()
	if bits != 32 || ones != 24 {
		die("Network %q must be a /24 (only /24 subnets are supported) — e.g. 10.0.0.0/24", network)
	}
	if ipnet.IP.String() != networkBase(network) {
		die("Network %q is not a valid /24 base address — did you mean %s/24?", network, ipnet.IP.String())
	}
}

// getFirstClientIP derives the .2 host of the configured network.
func getFirstClientIP(network string) string {
	base := networkBase(network)
	parts := strings.Split(base, ".")
	if len(parts) != 4 {
		die("Invalid network: %s", network)
	}
	return fmt.Sprintf("%s.%s.%s.2", parts[0], parts[1], parts[2])
}

// incrementIP returns the next host IP; if current is empty it seeds from .2.
func incrementIP(current, network string) string {
	if current == "" {
		return getFirstClientIP(network)
	}
	parts := strings.Split(current, ".")
	if len(parts) != 4 {
		die("Invalid IP: %s", current)
	}
	o4 := atoiOrDie(parts[3])
	o4++
	if o4 >= 255 {
		die("IP pool exhausted")
	}
	return fmt.Sprintf("%s.%s.%s.%d", parts[0], parts[1], parts[2], o4)
}

// ─── command execution ──────────────────────────────────────────────────────────

// run executes a command with inherited stdio and returns its error.
func run(name string, args ...string) error {
	return runEnv(nil, name, args...)
}

func runEnv(extraEnv []string, name string, args ...string) error {
	cmd := exec.Command(name, args...)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Stdin = os.Stdin
	if extraEnv != nil {
		cmd.Env = append(os.Environ(), extraEnv...)
	}
	return cmd.Run()
}

// runOrDie runs a command and aborts if it fails.
func runOrDie(name string, args ...string) {
	if err := run(name, args...); err != nil {
		die("%s failed: %v", name, err)
	}
}

// runOrDieIn runs a command with its working directory set to dir (e.g. a
// cloned source tree) and aborts if it fails.
func runOrDieIn(dir, name string, args ...string) {
	cmd := exec.Command(name, args...)
	cmd.Dir = dir
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Stdin = os.Stdin
	if err := cmd.Run(); err != nil {
		die("%s failed: %v", name, err)
	}
}

// enableIPForwarding turns on IPv4 forwarding and persists it. It is tolerant
// of container environments: if `sysctl -w` fails but forwarding is already on
// (e.g. set at namespace creation via `--sysctl`/compose, with /proc/sys mounted
// read-only) it warns and continues instead of aborting. It only dies if
// forwarding is genuinely off and cannot be enabled.
func enableIPForwarding() {
	if err := run("sysctl", "-w", "net.ipv4.ip_forward=1"); err != nil {
		if ipForwardingEnabled() {
			warn("Could not write net.ipv4.ip_forward (%v), but it is already enabled — continuing", err)
		} else {
			die("Failed to enable IP forwarding: %v (set it on the host, e.g. --sysctl net.ipv4.ip_forward=1)", err)
		}
	}
	// Best-effort persistence; /etc may be read-only in some containers.
	if err := os.WriteFile("/etc/sysctl.d/99-amneziawg.conf", []byte("net.ipv4.ip_forward=1\n"), 0644); err != nil {
		warn("Could not persist sysctl config: %v", err)
	}
	run("sysctl", "-p", "/etc/sysctl.d/99-amneziawg.conf")
}

// ipForwardingEnabled reads the live kernel flag directly.
func ipForwardingEnabled() bool {
	data, err := os.ReadFile("/proc/sys/net/ipv4/ip_forward")
	return err == nil && strings.TrimSpace(string(data)) == "1"
}

// output runs a command and returns its trimmed stdout.
func output(name string, args ...string) (string, error) {
	cmd := exec.Command(name, args...)
	cmd.Stderr = os.Stderr
	out, err := cmd.Output()
	return strings.TrimSpace(string(out)), err
}

// outputWithInput runs a command feeding stdin, returning trimmed stdout.
func outputWithInput(stdin, name string, args ...string) (string, error) {
	cmd := exec.Command(name, args...)
	cmd.Stdin = strings.NewReader(stdin)
	cmd.Stderr = os.Stderr
	out, err := cmd.Output()
	return strings.TrimSpace(string(out)), err
}

// silent reports whether a command succeeds, discarding all its output.
func silent(name string, args ...string) bool {
	cmd := exec.Command(name, args...)
	return cmd.Run() == nil
}

// ─── network detection ──────────────────────────────────────────────────────────

// defaultRouteIface parses `ip route` for the default outbound interface,
// mirroring `ip route | awk '/default/ {print $5; exit}'`.
func defaultRouteIface() string {
	out, err := output("ip", "route")
	if err != nil {
		return ""
	}
	for _, line := range strings.Split(out, "\n") {
		if strings.Contains(line, "default") {
			fields := strings.Fields(line)
			if len(fields) >= 5 {
				return fields[4]
			}
		}
	}
	return ""
}

// detectPublicIP mimics `curl -sf4 https://ifconfig.me` (IPv4 only), falling
// back to the placeholder used by the original script.
//
// It queries several plaintext "what is my IP" services in turn and validates
// each response with net.ParseIP. Validation is essential: services such as
// ifconfig.me serve a full HTML landing page (not the bare IP) to clients that
// don't send a curl-like User-Agent, so without the parse check the config
// would end up storing an HTML fragment instead of an address.
func detectPublicIP() string {
	dialer := &net.Dialer{Timeout: 5 * time.Second}
	transport := &http.Transport{
		DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
			// Force IPv4 so the detected address matches the udp4 listener.
			return dialer.DialContext(ctx, "tcp4", addr)
		},
	}
	client := &http.Client{Timeout: 8 * time.Second, Transport: transport}

	// Plaintext IPv4 endpoints, tried in order until one yields a valid IP.
	for _, url := range []string{
		"https://ifconfig.me/ip",
		"https://api.ipify.org",
		"https://icanhazip.com",
	} {
		if ip := fetchIPv4(client, url); ip != "" {
			return ip
		}
	}
	return "YOUR_SERVER_IP"
}

// fetchIPv4 requests url and returns the trimmed body only if it is a valid
// IPv4 address; otherwise it returns "". A curl-like User-Agent is sent so
// services that content-negotiate (e.g. ifconfig.me) reply with the bare IP.
func fetchIPv4(client *http.Client, url string) string {
	req, err := http.NewRequest(http.MethodGet, url, nil)
	if err != nil {
		return ""
	}
	req.Header.Set("User-Agent", "curl/8.0.0")
	req.Header.Set("Accept", "text/plain")

	resp, err := client.Do(req)
	if err != nil {
		return ""
	}
	defer resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return ""
	}
	// A valid IPv4 string is at most 15 bytes; cap the read to reject any
	// unexpectedly large (e.g. HTML) response early.
	body, err := io.ReadAll(io.LimitReader(resp.Body, 64))
	if err != nil {
		return ""
	}
	ip := strings.TrimSpace(string(body))
	if parsed := net.ParseIP(ip); parsed == nil || parsed.To4() == nil {
		return ""
	}
	return ip
}

// unameRelease returns the running kernel release (`uname -r`).
func unameRelease() string {
	out, err := output("uname", "-r")
	if err != nil {
		die("uname -r failed: %v", err)
	}
	return out
}