package main

import (
	"html"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"testing"
	"time"
)

func TestKeys(t *testing.T) {
	// RFC 7748's X25519 test vector: what a WireGuard public key is.
	pub, err := publicKey("dwdtCnMYpX08FsFyUbJmRd9ML4frwJkqsXf7pR25LCo=")
	if err != nil || pub != "hSDwCYkwp1R0i33ctD73Wg2/Og0mOBr066SpjqqbTmo=" {
		t.Fatalf("public key = %q %v", pub, err)
	}
	priv, pub, err := newPrivateKey()
	if err != nil || !b64KeyRe.MatchString(priv) || !b64KeyRe.MatchString(pub) {
		t.Fatalf("new key: %q %q %v", priv, pub, err)
	}
	if again, _ := publicKey(priv); again != pub {
		t.Error("the public key does not belong to the private key")
	}
	if _, err := publicKey("short"); err == nil {
		t.Error("a bad key was accepted")
	}
}

func TestParseWGDump(t *testing.T) {
	dump := "wg-a\tPRIV\tPUB\t51820\toff\n" +
		"wg-a\tPEERKEY=\t(none)\t203.0.113.7:51820\t10.0.0.2/32\t1700000000\t1234\t5678\t25\n" +
		"wg-a\tOTHER=\t(none)\t(none)\t10.0.0.3/32\t0\t0\t0\toff\n"
	got := parseWGDump(dump)["wg-a"]
	if p := got["PEERKEY="]; p.RX != 1234 || p.TX != 5678 || p.Endpoint != "203.0.113.7:51820" || !p.Handshake.Equal(time.Unix(1700000000, 0)) {
		t.Errorf("peer = %+v", p)
	}
	if p := got["OTHER="]; !p.Handshake.IsZero() || p.Endpoint != "" {
		t.Errorf("idle peer = %+v", p)
	}
}

const peerKey = "hSDwCYkwp1R0i33ctD73Wg2/Og0mOBr066SpjqqbTmo="

func TestApplyTunnel(t *testing.T) {
	sys := System{StateDir: t.TempDir()}
	s := mustParse(t)
	files, err := s.applyTunnel(sys, WGTunnel{
		Name: "family", ListenPort: "51820", Addresses: "10.100.0.1/24, fd42:100::1/64", DNS: true, KillSwitch: true,
		Peers: []WGPeer{
			{Name: "gran", PublicKey: peerKey, Allowed: "10.100.0.2/32", Keepalive: "25", NewPSK: true},
			{}, // the empty add row
		},
	})
	if err != nil {
		t.Fatal(err)
	}
	if len(files) != 2 || files[0].Path != sys.wgKeyPath("family") || files[1].Path != sys.wgPSKPath("family", "gran") {
		t.Fatalf("files = %+v", files)
	}
	if err := writeSecrets(sys.wgDir(), files); err != nil {
		t.Fatal(err)
	}
	if st, err := os.Stat(sys.wgKeyPath("family")); err != nil || st.Mode().Perm() != 0o600 {
		t.Errorf("key file: %v %v", st, err)
	}
	tunnels := s.tunnels(sys)
	if len(tunnels) != 2 || tunnels[0].Name != "family" || !tunnels[0].KeyPresent || tunnels[0].PublicKey == "" || !tunnels[0].DNS ||
		len(tunnels[0].Peers) != 1 || !tunnels[0].Peers[0].PSKSet {
		t.Fatalf("tunnels = %+v", tunnels)
	}
	// What else is in the settings survives (the sample has a wireguard "vpn" too).
	if got := str(s, "router", "wireguard", "vpn", "listenPort"); got != "51820" {
		t.Errorf("the other tunnel changed: %q", got)
	}
	out, _ := s.marshal()
	if strings.Contains(string(out), tunnels[0].PublicKey+"x") || strings.Contains(string(out), "PRIVATE") {
		t.Error("secrets in the settings")
	}

	// Editing keeps the key; removing the peer row removes the peer.
	before, _ := os.ReadFile(sys.wgKeyPath("family"))
	edit := tunnels[0]
	edit.Peers[0].Remove = true
	files, err = s.applyTunnel(sys, edit)
	if err != nil || len(files) != 0 {
		t.Fatalf("edit: %v %+v", err, files)
	}
	if after, _ := os.ReadFile(sys.wgKeyPath("family")); string(after) != string(before) {
		t.Error("the key changed")
	}
	if len(s.tunnels(sys)[0].Peers) != 0 {
		t.Error("peer not removed")
	}
	// A new key on request.
	edit.NewKey = true
	edit.Peers = nil
	files, _ = s.applyTunnel(sys, edit)
	if len(files) != 1 || files[0].Content == strings.TrimSpace(string(before)) {
		t.Errorf("no new key: %+v", files)
	}
	if !s.deleteTunnel("family") || s.deleteTunnel("family") {
		t.Error("delete")
	}
}

func TestApplyTunnelValidates(t *testing.T) {
	sys := System{StateDir: t.TempDir()}
	for name, tun := range map[string]WGTunnel{
		"name":               {Name: "Bad Name"},
		"listenPort":         {Name: "a", ListenPort: "70000"},
		"addresses":          {Name: "a", Addresses: "10.0.0.1"},
		"privateKey":         {Name: "a", PrivateKey: "nope"},
		"peers[0].publicKey": {Name: "a", Peers: []WGPeer{{Name: "p", PublicKey: "x", Allowed: "10.0.0.2/32"}}},
		"peers[0].allowed":   {Name: "a", Peers: []WGPeer{{Name: "p", PublicKey: peerKey, Allowed: "10.0.0.2"}}},
		"peers[0].endpoint":  {Name: "a", Peers: []WGPeer{{Name: "p", PublicKey: peerKey, Endpoint: "host", Allowed: "10.0.0.2/32"}}},
		"peers[1].publicKey": {Name: "a", Peers: []WGPeer{{Name: "p", PublicKey: peerKey}, {Name: "q", PublicKey: peerKey}}},
	} {
		s := mustParse(t)
		_, err := s.applyTunnel(sys, tun)
		fe, ok := err.(FormErrors)
		if !ok || fe[name] == "" {
			t.Errorf("%s: %v", name, err)
		}
	}
	// A pasted key is stored; an existing tunnel without a key file is refused.
	s := mustParse(t)
	files, err := s.applyTunnel(sys, WGTunnel{Name: "pasted", PrivateKey: "dwdtCnMYpX08FsFyUbJmRd9ML4frwJkqsXf7pR25LCo="})
	if err != nil || len(files) != 1 || files[0].Content != "dwdtCnMYpX08FsFyUbJmRd9ML4frwJkqsXf7pR25LCo=" {
		t.Fatalf("pasted: %v %+v", err, files)
	}
	if _, err := mustParse(t).applyTunnel(sys, WGTunnel{Name: "vpn"}); err == nil {
		t.Error("an existing tunnel without a key file was accepted")
	}
}

func TestNextClientAddress(t *testing.T) {
	tun := WGTunnel{Addresses: "10.100.0.1/24", Peers: []WGPeer{{Allowed: "10.100.0.2/32"}, {Allowed: "10.100.0.3/32"}}}
	if got := nextClientAddress(tun); got != "10.100.0.4" {
		t.Errorf("next = %q", got)
	}
	if got := nextClientAddress(WGTunnel{Addresses: "fd42::1/64"}); got != "" {
		t.Errorf("no IPv4 network: %q", got)
	}
}

func TestWireGuardPage(t *testing.T) {
	a, sys := newTestApp(t)
	h := a.routes()
	cookie := do(h, "POST", "/login", url.Values{"password": {"hunter2hunter2"}}, nil).Result().Cookies()[0]
	token := func(path string) string {
		body := do(h, "GET", path, nil, cookie).Body.String()
		return regexp.MustCompile(`name="csrf" value="([^"]+)"`).FindStringSubmatch(body)[1]
	}
	if body := do(h, "GET", "/wireguard", nil, cookie).Body.String(); !strings.Contains(body, "New tunnel") {
		t.Fatalf("no page:\n%s", body)
	}

	// A new tunnel gets its key; a bad one is refused with the form kept.
	form := url.Values{"csrf": {token("/wireguard")}, "name": {"family"}, "listenPort": {"51820"}, "addresses": {"10.100.0.1/24"},
		"killSwitch": {"on"}, "peers": {"1"}, "peer.0.name": {"gran"}, "peer.0.publicKey": {"bad"}, "peer.0.allowed": {"10.100.0.2/32"}}
	rec := do(h, "POST", "/wireguard/save", form, cookie)
	if rec.Code != http.StatusUnprocessableEntity || !strings.Contains(rec.Body.String(), `value="gran"`) {
		t.Fatalf("bad peer: %d", rec.Code)
	}
	if _, err := os.Stat(sys.wgKeyPath("family")); err == nil {
		t.Error("a key was written for a refused form")
	}
	form.Set("peer.0.publicKey", peerKey)
	if rec := do(h, "POST", "/wireguard/save", form, cookie); rec.Code != http.StatusSeeOther {
		t.Fatalf("save: %d\n%s", rec.Code, rec.Body)
	}
	if _, err := os.Stat(sys.wgKeyPath("family")); err != nil {
		t.Fatal("no key file")
	}
	s, _ := loadSettings(sys.settingsPath())
	if str(s, "router", "wireguard", "family", "privateKeyFile") != sys.wgKeyPath("family") {
		t.Fatalf("settings: %v", s)
	}
	body := html.UnescapeString(do(h, "GET", "/wireguard", nil, cookie).Body.String())
	pub := s.tunnels(sys)[0].PublicKey
	priv, _ := os.ReadFile(sys.wgKeyPath("family"))
	if !strings.Contains(body, pub) || strings.Contains(body, strings.TrimSpace(string(priv))) {
		t.Error("the page lacks the public key (or shows the private one)")
	}

	// A phone: the router makes its keys; the config is shown once.
	client := url.Values{"csrf": {token("/wireguard")}, "tunnel": {"family"}, "name": {"phone"}, "address": {"10.100.0.3"},
		"endpoint": {"vpn.example.org:51820"}, "allowed": {"10.100.0.0/24, 192.168.10.0/24"}, "dns": {"10.100.0.1"}}
	rec = do(h, "POST", "/wireguard/client", client, cookie)
	if rec.Code != http.StatusOK {
		t.Fatalf("client: %d\n%s", rec.Code, rec.Body)
	}
	shown := html.UnescapeString(rec.Body.String())
	for _, want := range []string{"[Interface]", "PrivateKey = ", "Address = 10.100.0.3/32", "DNS = 10.100.0.1", "PublicKey = " + pub, "PresharedKey = ", "Endpoint = vpn.example.org:51820"} {
		if !strings.Contains(shown, want) {
			t.Errorf("config lacks %q", want)
		}
	}
	s, _ = loadSettings(sys.settingsPath())
	peer := s.tunnels(sys)[0].Peers
	if len(peer) != 2 || peer[1].Name != "phone" || !peer[1].PSKSet || peer[1].Allowed != "10.100.0.3/32" {
		t.Fatalf("peers = %+v", peer)
	}
	psk, _ := os.ReadFile(sys.wgPSKPath("family", "phone"))
	if !strings.Contains(html.UnescapeString(rec.Body.String()), strings.TrimSpace(string(psk))) {
		t.Error("the config's preshared key is not the stored one")
	}
	if st, _ := os.Stat(filepath.Join(sys.wgDir(), "family-phone.psk")); st.Mode().Perm() != 0o600 {
		t.Error("psk permissions")
	}
	if _, err := exec.LookPath("qrencode"); err == nil && !strings.Contains(rec.Body.String(), "<svg") {
		t.Error("no QR code although qrencode is available")
	}
	// Not shown again, and the address is taken now.
	if body := do(h, "GET", "/wireguard", nil, cookie).Body.String(); strings.Contains(body, "Address = 10.100.0.3/32") {
		t.Error("the client config is shown again")
	}
	client.Set("csrf", token("/wireguard"))
	client.Set("name", "tablet")
	if rec := do(h, "POST", "/wireguard/client", client, cookie); rec.Code != http.StatusUnprocessableEntity {
		t.Errorf("a used address was accepted: %d", rec.Code)
	}

	// Delete needs the box ticked.
	do(h, "POST", "/wireguard/delete", url.Values{"csrf": {token("/wireguard")}, "name": {"family"}}, cookie)
	if s, _ := loadSettings(sys.settingsPath()); len(s.tunnels(sys)) != 2 { // vpn and family
		t.Error("deleted without confirmation")
	}
	do(h, "POST", "/wireguard/delete", url.Values{"csrf": {token("/wireguard")}, "name": {"family"}, "confirm": {"on"}}, cookie)
	if s, _ := loadSettings(sys.settingsPath()); len(s.tunnels(sys)) != 1 {
		t.Error("not deleted")
	}
	if _, err := os.Stat(sys.wgKeyPath("family")); err != nil {
		t.Error("the key file went with the tunnel")
	}
}

// With wg available: the keys made here are the keys WireGuard derives.
func TestKeysMatchWG(t *testing.T) {
	wg, err := exec.LookPath("wg")
	if err != nil {
		t.Skip("wg is not installed")
	}
	priv, pub, _ := newPrivateKey()
	cmd := exec.Command(wg, "pubkey")
	cmd.Stdin = strings.NewReader(priv + "\n")
	out, err := cmd.Output()
	if err != nil || strings.TrimSpace(string(out)) != pub {
		t.Fatalf("wg pubkey = %q %v, want %q", out, err, pub)
	}
}

const sampleConf = `# a comment before the interface
[Interface]
PrivateKey = dwdtCnMYpX08FsFyUbJmRd9ML4frwJkqsXf7pR25LCo=
Address = 10.100.0.1/24, fd42:100::1/64
ListenPort = 51820
MTU = 1380
DNS = 10.100.0.1
PostUp = iptables -A FORWARD -i %i -j ACCEPT

# Gran's laptop
[Peer]
PublicKey = hSDwCYkwp1R0i33ctD73Wg2/Og0mOBr066SpjqqbTmo=
PresharedKey = AAECAwQFBgcICQoLDA0ODxAREhMUFRYXGBkaGxwdHh8=
AllowedIPs = 10.100.0.2/32

[Peer]
PublicKey = 78NK8vKabw7TOKjZ4hjW/GgDLPFI03EcKksGUGuiC0I=
Endpoint = vpn.example.org:51820
AllowedIPs = 10.100.0.3/32, 192.168.50.0/24
PersistentKeepalive = 25
`

func TestParseWGQuick(t *testing.T) {
	imp, err := parseWGQuick("family", sampleConf)
	if err != nil {
		t.Fatal(err)
	}
	tun := imp.Tunnel
	if tun.PrivateKey == "" || tun.ListenPort != "51820" || tun.MTU != "1380" || tun.Addresses != "10.100.0.1/24, fd42:100::1/64" {
		t.Errorf("interface = %+v", tun)
	}
	if len(tun.Peers) != 2 || tun.Peers[0].Name != "gran-s-laptop" || tun.Peers[1].Name != "peer1" {
		t.Fatalf("peers = %+v", tun.Peers)
	}
	if p := tun.Peers[1]; p.Endpoint != "vpn.example.org:51820" || p.Keepalive != "25" || p.Allowed != "10.100.0.3/32, 192.168.50.0/24" {
		t.Errorf("peer = %+v", p)
	}
	if tun.Peers[0].PSK == "" {
		t.Error("the preshared key was lost")
	}
	if strings.Join(imp.Ignored, ",") != "dns,postup" && strings.Join(imp.Ignored, ",") != "postup,dns" {
		t.Errorf("ignored = %v", imp.Ignored)
	}
	// An address without a prefix length gets one; showconf has no Address at all.
	imp, err = parseWGQuick("x", "[Interface]\nPrivateKey = k\nAddress = 10.0.0.2\n")
	if err != nil || imp.Tunnel.Addresses != "10.0.0.2/32" {
		t.Errorf("address: %q %v", imp.Tunnel.Addresses, err)
	}

	for name, conf := range map[string]string{
		"no key":     "[Interface]\nAddress = 10.0.0.1/24\n",
		"full":       "[Interface]\nPrivateKey = k\n[Peer]\nPublicKey = p\nAllowedIPs = 0.0.0.0/0, ::/0\n",
		"garbage":    "[Interface]\nPrivateKey = k\nnonsense\n",
		"no section": "PrivateKey = k\n",
	} {
		if _, err := parseWGQuick("x", conf); err == nil {
			t.Errorf("%s: accepted", name)
		}
	}
}

func TestImportThroughThePage(t *testing.T) {
	a, sys := newTestApp(t)
	h := a.routes()
	cookie := do(h, "POST", "/login", url.Values{"password": {"hunter2hunter2"}}, nil).Result().Cookies()[0]
	token := func() string {
		body := do(h, "GET", "/wireguard", nil, cookie).Body.String()
		return regexp.MustCompile(`name="csrf" value="([^"]+)"`).FindStringSubmatch(body)[1]
	}

	// A full-tunnel peer is refused, and the key is not echoed back.
	bad := strings.Replace(sampleConf, "AllowedIPs = 10.100.0.2/32", "AllowedIPs = 0.0.0.0/0", 1)
	rec := do(h, "POST", "/wireguard/import", url.Values{"csrf": {token()}, "importName": {"family"}, "importText": {bad}}, cookie)
	if rec.Code != http.StatusUnprocessableEntity || !strings.Contains(rec.Body.String(), "routes by destination only") {
		t.Fatalf("full tunnel: %d", rec.Code)
	}
	for _, secret := range []string{"dwdtCnMYpX08FsFyUbJmRd9ML4frwJkqsXf7pR25LCo", "AAECAwQFBgcICQoLDA0ODxAREhMUFRYXGBkaGxwdHh8"} {
		if strings.Contains(html.UnescapeString(rec.Body.String()), secret) {
			t.Error("a key was echoed back")
		}
	}
	if _, err := os.Stat(sys.wgKeyPath("family")); err == nil {
		t.Error("a key was stored for a refused import")
	}

	rec = do(h, "POST", "/wireguard/import", url.Values{"csrf": {token()}, "importName": {"family"}, "importText": {sampleConf}}, cookie)
	if rec.Code != http.StatusSeeOther {
		t.Fatalf("import: %d\n%s", rec.Code, rec.Body)
	}
	if key, _ := os.ReadFile(sys.wgKeyPath("family")); strings.TrimSpace(string(key)) != "dwdtCnMYpX08FsFyUbJmRd9ML4frwJkqsXf7pR25LCo=" {
		t.Errorf("the imported key was not stored as it was: %q", key)
	}
	if st, _ := os.Stat(sys.wgPSKPath("family", "gran-s-laptop")); st == nil || st.Mode().Perm() != 0o600 {
		t.Error("the preshared key was not stored")
	}
	s, _ := loadSettings(sys.settingsPath())
	tun := s.tunnels(sys)[0]
	if tun.Name != "family" || tun.PublicKey != "hSDwCYkwp1R0i33ctD73Wg2/Og0mOBr066SpjqqbTmo=" || len(tun.Peers) != 2 || tun.Peers[1].Allowed != "10.100.0.3/32, 192.168.50.0/24" {
		t.Errorf("tunnel = %+v", tun)
	}
	if out, _ := s.marshal(); strings.Contains(string(out), "dwdtCnMY") || strings.Contains(string(out), "AAECAwQF") {
		t.Error("a secret is in settings.json")
	}
	if rec := do(h, "POST", "/wireguard/import", url.Values{"csrf": {token()}, "importName": {"family"}, "importText": {sampleConf}}, cookie); rec.Code != http.StatusUnprocessableEntity {
		t.Errorf("a second tunnel of the same name: %d", rec.Code)
	}
}