// godlighty -- highly-customizable HTTP, HTTP/2, HTTPS server
// Copyright (C) 2021-2024 Sergey Matveev <stargrave@stargrave.org>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, version 3 of the License.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package godlighty

import (
	"crypto/tls"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"log"
	"os"
)

var (
	NextProtos = []string{"h2", "http/1.1"}

	HostToECDSACertificate map[string]*tls.Certificate
	HostECDSAClientAuth    map[string]*x509.CertPool

	HostToEdDSACertificate map[string]*tls.Certificate
	HostEdDSAClientAuth    map[string]*x509.CertPool

	HostToGOSTCertificate map[string]*tls.Certificate
	HostGOSTClientAuth    map[string]*x509.CertPool
)

func CHIHasTLS13(chi *tls.ClientHelloInfo) bool {
	for _, v := range chi.SupportedVersions {
		if v == tls.VersionTLS13 {
			return true
		}
	}
	return false
}

func CHIHasEdDSA(chi *tls.ClientHelloInfo) bool {
	if !CHIHasTLS13(chi) {
		return false
	}
	for _, ss := range chi.SignatureSchemes {
		if ss == tls.Ed25519 {
			return true
		}
	}
	return false
}

func GetCertificate(chi *tls.ClientHelloInfo) (*tls.Certificate, error) {
	if CHIHasGOST(chi) {
		if cert := HostToGOSTCertificate[chi.ServerName]; cert != nil {
			return cert, nil
		}
	}
	var cert *tls.Certificate
	if len(HostToECDSACertificate) == 0 {
		cert = HostToEdDSACertificate[chi.ServerName]
	} else {
		if CHIHasEdDSA(chi) {
			if cert := HostToEdDSACertificate[chi.ServerName]; cert != nil {
				return cert, nil
			}
		}
		cert = HostToECDSACertificate[chi.ServerName]
	}
	if cert == nil {
		return nil, errors.New("no certificate found")
	}
	return cert, nil
}

func GetConfigForClient(chi *tls.ClientHelloInfo) (*tls.Config, error) {
	var pool *x509.CertPool
	if CHIHasGOST(chi) {
		pool = HostGOSTClientAuth[chi.ServerName]
	}
	if pool == nil && (CHIHasEdDSA(chi) || len(HostECDSAClientAuth) == 0) {
		pool = HostEdDSAClientAuth[chi.ServerName]
	}
	if pool == nil {
		pool = HostECDSAClientAuth[chi.ServerName]
	}
	if pool == nil {
		return nil, nil
	}
	return &tls.Config{
		GetCertificate: GetCertificate,
		NextProtos:     NextProtos,
		ClientCAs:      pool,
		ClientAuth:     tls.RequireAndVerifyClientCert,
	}, nil
}

func loadCertificates(
	host string,
	cfg *TLSCfg,
	hostToCertificate *map[string]*tls.Certificate,
	hostClientAuth *map[string]*x509.CertPool,
) {
	if cfg == nil {
		return
	}
	cert, err := tls.LoadX509KeyPair(cfg.Cert, cfg.Key)
	if err != nil {
		log.Fatalln(err)
	}
	if cfg.CACert != "" {
		data, err := os.ReadFile(cfg.CACert)
		if err != nil {
			log.Fatalln(err)
		}
		block, _ := pem.Decode(data)
		if block == nil {
			log.Fatalln(fmt.Errorf("no PEM found: %s", cfg.CACert))
		}
		if block.Type != "CERTIFICATE" {
			log.Fatalln(fmt.Errorf("non CERTIFICATE: %s", cfg.CACert))
		}
		cert.Certificate = append(cert.Certificate, block.Bytes)
	}
	(*hostToCertificate)[host] = &cert
	pool := x509.NewCertPool()
	for _, p := range cfg.ClientCAs {
		data, err := os.ReadFile(p)
		if err != nil {
			log.Fatalln(err)
		}
		var block *pem.Block
		for len(data) > 0 {
			block, data = pem.Decode(data)
			if block == nil {
				log.Fatalln("can not decode PEM:", p)
			}
			if block.Type != "CERTIFICATE" {
				continue
			}
			ca, err := x509.ParseCertificate(block.Bytes)
			if err != nil {
				log.Fatalln(err)
			}
			pool.AddCert(ca)
			(*hostClientAuth)[host] = pool
		}
	}
}

func LoadCertificates() {
	HostToECDSACertificate = make(map[string]*tls.Certificate, len(Hosts))
	HostECDSAClientAuth = make(map[string]*x509.CertPool)
	HostToEdDSACertificate = make(map[string]*tls.Certificate, len(Hosts))
	HostEdDSAClientAuth = make(map[string]*x509.CertPool)
	HostToGOSTCertificate = make(map[string]*tls.Certificate, len(Hosts))
	HostGOSTClientAuth = make(map[string]*x509.CertPool)
	for host, cfg := range Hosts {
		loadCertificates(host, cfg.ECDSATLS, &HostToECDSACertificate, &HostECDSAClientAuth)
		loadCertificates(host, cfg.EdDSATLS, &HostToEdDSACertificate, &HostEdDSAClientAuth)
		loadCertificates(host, cfg.GOSTTLS, &HostToGOSTCertificate, &HostGOSTClientAuth)
	}
}

func NewTLSConfig() *tls.Config {
	return &tls.Config{
		MinVersion:         tls.VersionTLS12,
		NextProtos:         NextProtos,
		GetCertificate:     GetCertificate,
		GetConfigForClient: GetConfigForClient,
	}
}