/* SPDX-License-Identifier: GPL-2.0
 *
 * Copyright (C) 2021 Jason A. Donenfeld. All Rights Reserved.
 */

package main

import (
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"regexp"
	"strings"

	"golang.org/x/crypto/curve25519"
	"inet.af/netaddr"
)

const vpnProviderMaxEndpointsInParallel = 20

type vpnEndpoint struct {
	publicKey []byte
	endpoint  netaddr.IPPort
	name      string
	country   string
}

type vpnConf struct {
	privateKey []byte
	ips        []netaddr.IP
	dnses      []netaddr.IP
}

type apiRelaysWireGuardV1Key []byte

func (key *apiRelaysWireGuardV1Key) UnmarshalText(text []byte) error {
	k, err := base64.StdEncoding.DecodeString(string(text))
	if err != nil {
		return err
	}
	if len(k) != 32 {
		return errors.New("key must be 32 bytes")
	}
	*key = k
	return nil
}

type apiRelaysWireGuardV1Relay struct {
	Hostname     string                  `json:"hostname"`
	EndpointV4   netaddr.IP              `json:"ipv4_addr_in"`
	EndpointV6   netaddr.IP              `json:"ipv6_addr_in"`
	PublicKey    apiRelaysWireGuardV1Key `json:"public_key"`
	MultihopPort uint16                  `json:"multihop_port"`
}

type apiRelaysWireGuardV1City struct {
	Name      string                      `json:"name"`
	Code      string                      `json:"code"`
	Latitude  float64                     `json:"latitude"`
	Longitude float64                     `json:"Longitude"`
	Relays    []apiRelaysWireGuardV1Relay `json:"relays"`
}

type apiRelaysWireGuardV1Country struct {
	Name   string                     `json:"name"`
	Code   string                     `json:"code"`
	Cities []apiRelaysWireGuardV1City `json:"cities"`
}
type apiRelaysWireGuardV1Root struct {
	Countries []apiRelaysWireGuardV1Country `json:"countries"`
}

func getVpnEndpoints() ([]vpnEndpoint, error) {
	log.Println("Getting VPN server list")
	resp, err := http.Get("https://api.mullvad.net/public/relays/wireguard/v1/")
	if err != nil {
		return nil, err
	}
	bytes, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	var relays apiRelaysWireGuardV1Root
	err = json.Unmarshal(bytes, &relays)
	if err != nil {
		return nil, err
	}
	var endpoints []vpnEndpoint
	for _, country := range relays.Countries {
		for _, city := range country.Cities {
			for _, relay := range city.Relays {
				endpoints = append(endpoints, vpnEndpoint{
					publicKey: relay.PublicKey,
					endpoint:  netaddr.IPPortFrom(relay.EndpointV4, 51820),
					name:      strings.TrimSuffix(relay.Hostname, "-wireguard"),
				})
			}
		}
	}
	return endpoints, nil
}

func registerVpnConf(privateKey []byte, accountId string) (*vpnConf, error) {
	log.Println("Registering VPN private key")
	publicKey, err := curve25519.X25519(privateKey, curve25519.Basepoint)
	if err != nil {
		return nil, err
	}
	resp, err := http.PostForm("https://api.mullvad.net/wg/", url.Values{
		"account": {accountId},
		"pubkey":  {base64.StdEncoding.EncodeToString(publicKey)},
	})
	if err != nil {
		return nil, err
	}
	ipBytes, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	if match, _ := regexp.Match(`^[0-9a-f:/.,]+$`, ipBytes); !match {
		return nil, fmt.Errorf("registration rejected: %q", string(ipBytes))
	}
	conf := &vpnConf{privateKey: privateKey, dnses: []netaddr.IP{netaddr.MustParseIP("193.138.218.74")}}
	for _, ipStr := range strings.Split(string(ipBytes), ",") {
		ip, err := netaddr.ParseIPPrefix(strings.TrimSpace(ipStr))
		if err != nil {
			return nil, err
		}
		conf.ips = append(conf.ips, ip.IP())
	}
	return conf, nil
}