package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"image/png"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"strconv"
	"strings"

	"github.com/boombuler/barcode"
	"github.com/boombuler/barcode/qr"
	"github.com/urfave/cli/v2"
)

func generateKeys() (string, string) {
	cmd := exec.Command("wg", "genkey")
	var out1 bytes.Buffer
	cmd.Stdout = &out1
	cmd.Run()

	private := strings.ReplaceAll(out1.String(), "\n", "")

	cmd2 := exec.Command("wg", "pubkey")
	cmd2.Stdin = strings.NewReader(private)

	var out2 bytes.Buffer
	cmd2.Stdout = &out2

	cmd2.Run()
	public := strings.ReplaceAll(out2.String(), "\n", "")

	return private, public
}

func generateConfigs(storage *Storage) {
	os.MkdirAll("out/server", 0755)
	os.MkdirAll("out/clients", 0755)

	serverConfig := ""
	serverConfig += "[Interface]\n"
	serverConfig += fmt.Sprintf("ListenPort = %d\n", storage.Port)
	serverConfig += fmt.Sprintf("Address = %s1/32,%s1/128\n", storage.Subnet, storage.Subnet6)
	serverConfig += fmt.Sprintf("PrivateKey = %s\n", storage.Server.Private)

	if len(storage.Server.IPTablesUp) > 0 {
		serverConfig += fmt.Sprintf("PostUp = %s\n", strings.Join(storage.Server.IPTablesUp, "; "))
	}

	if len(storage.Server.IPTablesDown) > 0 {
		serverConfig += fmt.Sprintf("PostDown = %s\n", strings.Join(storage.Server.IPTablesDown, "; "))
	}

	serverConfig += "# ---------------------------\n\n"

	for name, peer := range storage.Clients {
		serverConfig += fmt.Sprintf("[Peer] # %s\n", name)
		serverConfig += fmt.Sprintf("PublicKey = %s\n", peer.Public)
		serverConfig += fmt.Sprintf("AllowedIPs = %s%d/32,%s%d/128\n", storage.Subnet, peer.IP, storage.Subnet6, peer.IP)

		// serverConfig += fmt.Sprintf("Endpoint = %s:%d\n", storage.Server.EndPoint, storage.Port)
		serverConfig += "\n"
	}

	os.WriteFile("out/server/server.conf", []byte(serverConfig), 0644)

	for name, peer := range storage.Clients {
		peerConfig := "[Interface]\n"
		peerConfig += fmt.Sprintf("Address = %s%d/32,%s%d/128\n", storage.Subnet, peer.IP, storage.Subnet6, peer.IP)
		peerConfig += fmt.Sprintf("PrivateKey = %s\n", peer.Private)
		peerConfig += fmt.Sprintf("DNS = 1.1.1.1\n")
		peerConfig += "\n"
		peerConfig += fmt.Sprintf("[Peer] # Server\n")
		peerConfig += fmt.Sprintf("PublicKey = %s\n", storage.Server.Public)
		if !peer.OnlySubnet {
			peerConfig += fmt.Sprintf("AllowedIPs = 0.0.0.0/0, ::0/0\n")
		} else {
			peerConfig += fmt.Sprintf("AllowedIPs = %s0/24, %s0/68\n", storage.Subnet, storage.Subnet6)
		}
		peerConfig += fmt.Sprintf("Endpoint = %s:%d\n", storage.Server.EndPoint, storage.Port)

		os.WriteFile(fmt.Sprintf("out/clients/%s.conf", name), []byte(peerConfig), 0644)
		qrCode, _ := qr.Encode(peerConfig, qr.L, qr.Auto)
		qrCode, _ = barcode.Scale(qrCode, 512, 512)

		file, err := os.Create(fmt.Sprintf("out/clients/%s.png", name))
		if err != nil {
			log.Fatal(err)
		}

		defer file.Close()

		png.Encode(file, qrCode)

		file.Sync()
	}
}

func initConfig(storage *Storage) {
	var text string

	scanner := bufio.NewScanner(os.Stdin)
	fmt.Printf("Enter IPv4 subnet (%s): ", storage.Subnet)
	scanner.Scan()
	text = scanner.Text()
	if len(text) > 0 {
		storage.Subnet = text
	}

	fmt.Printf("Enter IPv6 subnet (%s): ", storage.Subnet6)
	scanner.Scan()
	text = scanner.Text()
	if len(text) > 0 {
		storage.Subnet6 = text
	}

	fmt.Printf("Enter Server Endpoint (%s): ", storage.Server.EndPoint)
	scanner.Scan()
	text = scanner.Text()
	if len(text) > 0 {
		storage.Server.EndPoint = text
	}

	fmt.Printf("Enter Port (%d): ", storage.Port)
	scanner.Scan()
	text = scanner.Text()
	if len(text) > 0 {
		u, err := strconv.ParseUint(text, 10, 32)
		if err != nil {
			log.Fatal(err)
		}
		storage.Port = uint32(u)
	}
}

func main() {
	if !fileExists("config.json") {
		err := ioutil.WriteFile("config.json", []byte("{}"), os.FileMode(0750))
		if err != nil {
			fmt.Print(err)
			os.Exit(1)
		}
	}

	data, err := ioutil.ReadFile("config.json")
	if err != nil {
		log.Fatal(err)
	}

	storage := new(Storage)

	err = json.Unmarshal(data, &storage)

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

	if storage.Server.Private == "" || storage.Server.Public == "" {
		pr, pu := generateKeys()
		storage.Server.Private = pr
		storage.Server.Public = pu
	}

	for storage.Subnet == "" || storage.Subnet6 == "" || storage.Server.EndPoint == "" {
		initConfig(storage)
	}

	if storage.Clients == nil {
		storage.Clients = map[string]Client{}
	}

	if storage.Server.IPTablesUp == nil {
		storage.Server.IPTablesUp = []string{}
	}

	if storage.Server.IPTablesDown == nil {
		storage.Server.IPTablesDown = []string{}
	}

	// fmt.Printf("%+v\n", storage)

	app := &cli.App{
		Name:    "wgconf",
		Usage:   "make an explosive entrance",
		Version: "0.0.1",
		Commands: []*cli.Command{
			{
				Name:    "generate",
				Aliases: []string{"g"},
				Usage:   "Generate the wireguard configs",
				Action: func(c *cli.Context) error {
					generateConfigs(storage)
					return nil
				},
			}, {
				Name:    "add",
				Aliases: []string{"a"},
				Usage:   "Add device",
				Flags: []cli.Flag{
					&cli.BoolFlag{
						Name:    "onlysubnet",
						Aliases: []string{"s"},
						Usage:   "Only VPN subnet, not whole traffic",
						Value:   false,
					},
				},
				Action: func(c *cli.Context) error {
					name := c.Args().First()
					if len(name) <= 0 {
						log.Fatalln("Name not set!")
					}

					freeIP := 1

					for _, peer := range storage.Clients {
						if peer.IP >= freeIP {
							freeIP = peer.IP + 1
						}
					}

					private, public := generateKeys()

					storage.Clients[name] = Client{
						OnlySubnet: c.Bool("onlysubnet"),
						Private:    private,
						Public:     public,
						IP:         freeIP,
					}

					return nil
				},
			},
		},
		Action: func(c *cli.Context) error {
			fmt.Println("boom! I say!")
			return nil
		},
	}

	err = app.Run(os.Args)

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

	data, err = json.MarshalIndent(storage, "", "   ")

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

	err = os.WriteFile("config.json", data, 0644)
	if err != nil {
		log.Fatal(err)
	}
}

func fileExists(filename string) bool {
	info, err := os.Stat(filename)
	if os.IsNotExist(err) {
		return false
	}
	return !info.IsDir()
}

type Client struct {
	IP         int    `json:"ip"`
	Private    string `json:"private"`
	Public     string `json:"public"`
	OnlySubnet bool   `json:"onlysubnet"`
}

type Server struct {
	EndPoint     string   `json:"endpoint"`
	Private      string   `json:"private"`
	Public       string   `json:"public"`
	IPTablesUp   []string `json:"iptables-up"`
	IPTablesDown []string `json:"iptables-down"`
}

type Storage struct {
	Subnet  string            `json:"subnet"`
	Port    uint32            `json:"port"`
	Server  Server            `json:"server"`
	Clients map[string]Client `json:"clients"`
	Subnet6 string            `json:"subnet6"`
}