/*
tofuproxy -- HTTP proxy with TLS certificates management
Copyright (C) 2021 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 tofuproxy

import (
	"bytes"
	"crypto/sha256"
	"crypto/x509"
	"encoding/hex"
	"encoding/pem"
	"fmt"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"

	"go.cypherpunks.ru/ucspi"
	"go.stargrave.org/tofuproxy/caches"
	"go.stargrave.org/tofuproxy/fifos"
)

var (
	CmdCerttool = "certtool"
	CmdWish     = "wish8.6"

	Certs   string
	VerifyM sync.Mutex
)

func spkiHash(cert *x509.Certificate) string {
	hsh := sha256.Sum256(cert.RawSubjectPublicKeyInfo)
	return hex.EncodeToString(hsh[:])
}

type ErrRejected struct {
	addr string
}

func (err ErrRejected) Error() string { return err.addr + " was rejected" }

func certInfo(certRaw []byte) string {
	cmd := exec.Command(CmdCerttool, "--certificate-info", "--inder")
	cmd.Stdin = bytes.NewReader(certRaw)
	out, err := cmd.Output()
	if err != nil {
		return err.Error()
	}
	lines := make([]string, 0, 128)
	for i, line := range strings.Split(string(out), "\n") {
		if strings.Contains(line, "ASCII:") {
			continue
		}
		lines = append(lines, fmt.Sprintf(
			"%03d %s", i, strings.ReplaceAll(line, `"`, `\"`),
		))
	}
	return strings.Join(lines, "\n")
}

func verifyCert(
	host string,
	dialErr error,
	rawCerts [][]byte,
	verifiedChains [][]*x509.Certificate,
) error {
	var certTheir *x509.Certificate
	var err error
	if len(verifiedChains) > 0 {
		certTheir = verifiedChains[0][0]
	} else {
		certTheir, err = x509.ParseCertificate(rawCerts[0])
		if err != nil {
			return err
		}
	}
	certTheirHash := spkiHash(certTheir)
	VerifyM.Lock()
	defer VerifyM.Unlock()
	caches.AcceptedM.RLock()
	certOurHash := caches.Accepted[host]
	caches.AcceptedM.RUnlock()
	if certTheirHash == certOurHash {
		return nil
	}
	caches.RejectedM.RLock()
	certOurHash = caches.Rejected[host]
	caches.RejectedM.RUnlock()
	if certTheirHash == certOurHash {
		return ErrRejected{host}
	}
	daneExists, daneMatched := dane(host, certTheir)
	if daneExists {
		if daneMatched {
			fifos.LogDANE <- fmt.Sprintf("%s\tACK", host)
		} else {
			fifos.LogDANE <- fmt.Sprintf("%s\tNAK", host)
		}
	}
	fn := filepath.Join(Certs, host)
	certsOur, _, err := ucspi.CertPoolFromFile(fn)
	if err == nil || dialErr != nil || (daneExists && !daneMatched) {
		if certsOur != nil && certTheirHash == spkiHash(certsOur[0]) {
			caches.AcceptedM.Lock()
			caches.Accepted[host] = certTheirHash
			caches.AcceptedM.Unlock()
			if bytes.Compare(certsOur[0].Raw, rawCerts[0]) != 0 {
				fifos.LogCert <- fmt.Sprintf("Refresh\t%s\t%s", host, certTheirHash)
				goto CertUpdate
			}
			return nil
		}
		var b bytes.Buffer
		b.WriteString("tk_setPalette grey\n")
		b.WriteString(fmt.Sprintf("wm title . \"%s\"\n", host))

		if dialErr != nil {
			b.WriteString(fmt.Sprintf(`set tErr [text .tErr]
$tErr insert end "%s"
$tErr configure -wrap word -height 5
`, dialErr.Error()))
			b.WriteString("grid .tErr -columnspan 3\n")
		}

		if daneExists {
			if daneMatched {
				b.WriteString("label .lDANE -bg green -text \"DANE matched\"\n")
			} else {
				b.WriteString("label .lDANE -bg red -text \"DANE NOT matched\"\n")
			}
			b.WriteString("grid .lDANE\n")
		}

		var bCerts bytes.Buffer
		for i, rawCert := range rawCerts {
			bCerts.WriteString(fmt.Sprintf("Their %d:\n", i))
			bCerts.WriteString(certInfo(rawCert))
		}
		b.WriteString(fmt.Sprintf(`set tTheir [text .tTheir]
$tTheir insert end "%s"
set sbTheir [scrollbar .sbTheir -command [list $tTheir yview]]
$tTheir configure -wrap word -yscrollcommand [list $sbTheir set]
`, bCerts.String()))
		b.WriteString("grid $tTheir $sbTheir -sticky nsew -columnspan 3\n")

		if certsOur != nil {
			bCerts.Reset()
			for i, cert := range certsOur {
				bCerts.WriteString(fmt.Sprintf("Our %d:\n", i))
				bCerts.WriteString(certInfo(cert.Raw))
			}
			b.WriteString(fmt.Sprintf(`set tOur [text .tOur]
$tOur insert end "%s"
set sbOur [scrollbar .sbOur -command [list $tOur yview]]
$tOur configure -wrap word -yscrollcommand [list $sbOur set]
`, bCerts.String()))
			b.WriteString("grid $tOur $sbOur -sticky nsew -columnspan 3\n")
		}

		b.WriteString(`
proc doAccept {} { exit 10 }
proc doOnce {} { exit 11 }
proc doReject {} { exit 12 }
button .bAccept -text "Accept" -bg green -command doAccept
button .bOnce -text "Once" -bg green -command doOnce
button .bReject -text "Reject" -bg red -command doReject
grid .bAccept .bOnce .bReject
grid rowconfigure . 0 -weight 1
grid columnconfigure . 0 -weight 1
`)

		cmd := exec.Command(CmdWish)
		// ioutil.WriteFile("/tmp/w.tcl", b.Bytes(), 0666)
		cmd.Stdin = &b
		err = cmd.Run()
		exitError, ok := err.(*exec.ExitError)
		if !ok {
			fifos.LogCert <- fmt.Sprintf("Reject\t%s\t%s", host, certTheirHash)
			return ErrRejected{host}
		}
		switch exitError.ExitCode() {
		case 10:
			fifos.LogCert <- fmt.Sprintf("Accept\t%s\t%s", host, certTheirHash)
			goto CertUpdate
		case 11:
			fifos.LogCert <- fmt.Sprintf("Once\t%s\t%s", host, certTheirHash)
			caches.AcceptedM.Lock()
			caches.Accepted[host] = certTheirHash
			caches.AcceptedM.Unlock()
			return nil
		case 12:
			caches.RejectedM.Lock()
			caches.Rejected[host] = certTheirHash
			caches.RejectedM.Unlock()
			fallthrough
		default:
			fifos.LogCert <- fmt.Sprintf("Reject\t%s\t%s", host, certTheirHash)
			return ErrRejected{host}
		}
	} else {
		if !os.IsNotExist(err) {
			return err
		}
		fifos.LogCert <- fmt.Sprintf("TOFU\t%s\t%s", host, certTheirHash)
	}
CertUpdate:
	tmp, err := os.CreateTemp(Certs, "")
	if err != nil {
		log.Fatalln(err)
	}
	for _, rawCert := range rawCerts {
		err = pem.Encode(tmp, &pem.Block{Type: "CERTIFICATE", Bytes: rawCert})
		if err != nil {
			log.Fatalln(err)
		}
	}
	tmp.Close()
	os.Rename(tmp.Name(), fn)
	caches.AcceptedM.Lock()
	caches.Accepted[host] = certTheirHash
	caches.AcceptedM.Unlock()
	return nil
}