package main

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"os"
	"os/exec"
	"strconv"
	"strings"
)

const (
	IFaceName  = "wg0"
	TunnelPort = 12912
)

var (
	SrvPubKey string
)

func init() {
	// Generate private key
	outPrvK, err := exec.Command("wg", "genkey").Output()
	if err != nil {
		log.Fatal("Unable to generate prvkey:", err)
		return
	}
	outPrvK = bytes.TrimSpace(outPrvK)

	prvFile, err := ioutil.TempFile("", "wg")
	if err != nil {
		log.Fatal("Unable to create tmpPrvFile", err)
		return
	}
	defer os.Remove(prvFile.Name())

	if _, err := prvFile.Write(outPrvK); err != nil {
		log.Fatal(err)
	}
	if err := prvFile.Close(); err != nil {
		log.Fatal(err)
	}

	// Calculate public key
	cmdPubK := exec.Command("wg", "pubkey")
	cmdPubK.Stdin = bytes.NewReader(outPrvK)
	outPubK, err := cmdPubK.Output()
	if err != nil {
		log.Fatal("Unable to calculate pubkey:", err)
		return
	}
	SrvPubKey = string(bytes.TrimSpace(outPubK))

	// Set parameters
	outSet, err := exec.Command("wg", "set", IFaceName, "private-key", prvFile.Name(), "listen-port", strconv.Itoa(TunnelPort)).CombinedOutput()
	if err != nil {
		log.Fatal("Something happend during wg set:", err, string(outSet))
		return
	}

	log.Printf("wg configured: public key:%s", string(outPubK))
}

type PubTunnel struct {
	Login  string   `json:"login"`
	PubKey [][]byte `json:"data"`
	Token  string   `json:"token"`
}

func register(w http.ResponseWriter, r *http.Request) {
	if addr := r.Header.Get("X-Forwarded-For"); addr != "" {
		r.RemoteAddr = addr
	}
	log.Printf("%s \"%s %s\" [%s]\n", r.RemoteAddr, r.Method, r.URL.Path, r.UserAgent())

	// Read the body
	if r.ContentLength < 0 || r.ContentLength > 6553600 {
		http.Error(w, "{errmsg:\"Request too large or request size unknown\"}", http.StatusRequestEntityTooLarge)
		return
	}

	var pt PubTunnel
	if err := json.NewDecoder(r.Body).Decode(&pt); err != nil {
		http.Error(w, fmt.Sprintf("{errmsg:%q}", err), http.StatusBadRequest)
		return
	}

	// Validate wg token
	if j, err := json.Marshal(pt); err != nil {
		http.Error(w, fmt.Sprintf("{errmsg:%q}", err), http.StatusInternalServerError)
		return
	} else if r, err := http.NewRequest("POST", "https://adlin.nemunai.re/wg-step", bytes.NewReader(j)); err != nil {
		http.Error(w, fmt.Sprintf("{errmsg:%q}", err), http.StatusInternalServerError)
		return
	} else if resp, err := http.DefaultClient.Do(r); err == nil {
		resp.Body.Close()
	} else {
		log.Printf("Unable to register wg-step on token-validator:", err)
	}

	if next_ip, err := findNextIP(); err != nil {
		http.Error(w, fmt.Sprintf("{errmsg:%q}", err), http.StatusBadRequest)
		return
	} else {
		addWgPeer(pt.PubKey[0], next_ip)

		w.Header().Set("Content-Type", "text/plain")
		w.Write([]byte(fmt.Sprintf(`# Address=%s/18

[Peer]
PublicKey = %s
Endpoint = %s:%d
AllowedIPs = 0.0.0.0/0
PersistentKeepalive = 15
`, next_ip, SrvPubKey, "172.17.0.15", TunnelPort)))
	}
}

func findNextIP() (next_ip net.IP, err error) {
	var out []byte
	out, err = exec.Command("wg", "show", IFaceName, "allowed-ips").Output()
	if err != nil {
		return
	}

	all_ips := string(out)

	for j := byte(128); j < 190; j++ {
		for i := byte(1); i < 254; i++ {
			if !strings.Contains(all_ips, fmt.Sprintf("172.23.%d.%d/32", j, i)) {
				next_ip = net.IPv4(172, 23, j, i)
				return
			}
		}
	}

	err = errors.New("No IP found")
	return
}

func addWgPeer(cltPubKey []byte, next_ip net.IP) (err error) {
	_, err = exec.Command(
		"wg", "set", IFaceName,
		"peer", base64.StdEncoding.EncodeToString(cltPubKey),
		"allowed-ips", fmt.Sprintf("%s/32", next_ip),
	).Output()
	return
}