// tofuproxy -- flexible HTTP/HTTPS proxy, TLS terminator, X.509 TOFU
//              manager, WARC/geminispace browser
// 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 tofuproxy

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

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

const VerifyDialog = `
# host err daneStatus certsTheir certsOur

if {[string length $err] > 0} {
    set tErr [text .tErr]
    $tErr insert end $err
    $tErr configure -wrap word -height 5
    $tErr configure -state disabled
    grid .tErr
}

proc certsDecode {raws} {
    set certs [list]
    foreach raw [split $raws] {
        set lines [list]
        set lineN 0
        foreach line [split [binary decode hex $raw] "\n"] {
            lappend lines [format "%03d %s" $lineN $line]
            incr lineN
        }
        lappend certs [join $lines "\n"]
    }
    return $certs
}

set certsTheir [certsDecode $certsTheir]
set certsOur [certsDecode $certsOur]

tk_setPalette grey
wm title . $host

proc paginator {i delta t l certs} {
    incr i $delta
    if {$i == [llength $certs]} {
        set i 0
    } elseif {$i < 0} {
        set i [expr {[llength $certs] - 1}]
    }
    $t configure -state normal
    $t delete 1.0 end
    $t insert end [lindex $certs $i]
    $t configure -state disabled
    $l configure -text "[expr {$i + 1}] / [llength $certs]"
    return $i
}

proc addCertsWindow {name} {
    global certs$name t$name sb$name page$name l$name
    set t [text .t$name]
    set sb [scrollbar .sb$name -command [list $t yview]]
    $t configure -wrap word -yscrollcommand [list $sb set]
    $t configure -state disabled
    grid $t $sb -sticky nsew
    set t$name $t
    set sb$name $t

    frame .fControl$name
    set l$name [label .lPage$name]
    button .bNext$name -text "Next" -command [subst {
        set page$name \[paginator \$page$name +1 \$t$name \$l$name \$certs$name]
    }]
    button .bPrev$name -text "Prev" -command [subst {
        set page$name \[paginator \$page$name -1 \$t$name \$l$name \$certs$name]
    }]
    grid .fControl$name
    grid .lPage$name .bNext$name .bPrev$name -in .fControl$name
    set page$name [paginator -1 +1 $t [set l$name] [set certs$name]]
}

addCertsWindow Their
if {[llength $certsOur] > 0} { addCertsWindow Our }
frame .fButtons
set lDANE [label .lDANE]
if {$daneStatus ne ""} {
    array set daneColour {ok green bad red}
    $lDANE configure -bg $daneColour($daneStatus)
    $lDANE configure -text "DANE-EE: $daneStatus"
}
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 .fButtons
grid .lDANE .bAccept .bOnce .bReject -in .fButtons
grid rowconfigure . 0 -weight 1
grid columnconfigure . 0 -weight 1

bind . <KeyPress> {switch -exact %K {
    q {exit 0} ; # reject once
    a doAccept
    o doOnce
    r doReject
    n {.bNextTheir invoke}
    p {.bPrevTheir invoke}
    N {.bNextOur invoke}
    P {.bPrevOur invoke}
}}
`

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()
	}
	return string(out)
}

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)
		}
	}
	if len(verifiedChains) > 0 {
		caHashes := make(map[string]struct{})
		for _, certs := range verifiedChains {
			for _, cert := range certs {
				caHashes[spkiHash(cert)] = struct{}{}
			}
		}
		var restrictedHosts []string
		caches.RestrictedM.RLock()
		for h := range caHashes {
			restrictedHosts = append(restrictedHosts, caches.Restricted[h]...)
		}
		caches.RestrictedM.RUnlock()
		if len(restrictedHosts) > 0 {
			for _, h := range restrictedHosts {
				if host == h || strings.HasSuffix(host, "."+h) {
					goto HostIsNotRestricted
				}
			}
			fifos.LogCert <- fmt.Sprintf("Restricted\t%s", host)
			return ErrRejected{host}
		}
	}
HostIsNotRestricted:
	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.Equal(certsOur[0].Raw, rawCerts[0]) {
				fifos.LogCert <- fmt.Sprintf("Refresh\t%s\t%s", host, certTheirHash)
				goto CertUpdate
			}
			return nil
		}
		var b bytes.Buffer
		fmt.Fprintf(&b, "set host \"%s\"\n", host)
		if dialErr == nil {
			fmt.Fprintf(&b, "set err \"\"\n")
		} else {
			fmt.Fprintf(&b, "set err \"%s\"\n", dialErr.Error())
		}
		var daneStatus string
		if daneExists {
			if daneMatched {
				daneStatus = "ok"
			} else {
				daneStatus = "bad"
			}
		}
		fmt.Fprintf(&b, "set daneStatus \"%s\"\n", daneStatus)
		hexCerts := make([]string, 0, len(rawCerts))
		for _, rawCert := range rawCerts {
			hexCerts = append(hexCerts, hex.EncodeToString([]byte(certInfo(rawCert))))
		}
		fmt.Fprintf(&b, "set certsTheir \"%s\"\n", strings.Join(hexCerts, " "))
		hexCerts = make([]string, 0, len(certsOur))
		for _, cert := range certsOur {
			hexCerts = append(hexCerts, hex.EncodeToString([]byte(certInfo(cert.Raw))))
		}
		fmt.Fprintf(&b, "set certsOur \"%s\"\n", strings.Join(hexCerts, " "))
		b.WriteString(VerifyDialog)
		cmd := exec.Command(CmdWish)
		// os.WriteFile("/tmp/verify-dialog.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 !errors.Is(err, fs.ErrNotExist) {
			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
}