/*
ucspi/cmd/tlsc -- UCSPI TLS client
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 main

import (
	"bytes"
	"crypto/sha256"
	"crypto/tls"
	"crypto/x509"
	"encoding/hex"
	"encoding/pem"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"

	"go.cypherpunks.ru/ucspi"
)

func main() {
	onlyShow := flag.Bool("show", false, "Show server's certificate and exit")
	crtPath := flag.String("cert", "", "Path to client X.509 certificate")
	prvPath := flag.String("key", "", "Path to client PKCS#8 private key")
	casPath := flag.String("ca", "", "Path to CA certificates file")
	hostname := flag.String("name", "example.com", "Expected server's hostname")
	insecure := flag.Bool("insecure", false, "Insecure mode")
	fpr := flag.String("fpr", "", "Expected SHA256 hash of server certificate's SPKI")
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, `Usage: tcpclient host port tlsc -name expected.name
        [-cert cert.pem -key prv.pem] [-ca CAs.pem] [-fpr DEADBABE] [-show]
        program [args...]

`)
		flag.PrintDefaults()
	}
	flag.Parse()
	log.SetFlags(log.Lshortfile)

	if len(flag.Args()) == 0 {
		flag.Usage()
		os.Exit(1)
	}

	cfg := &tls.Config{}
	if *hostname == "" || *onlyShow || *insecure {
		cfg.InsecureSkipVerify = true
	}
	if *hostname != "" {
		cfg.ServerName = *hostname
	}
	if *crtPath != "" {
		crtRaw, _, err := ucspi.CertificateFromFile(*crtPath)
		if err != nil {
			log.Fatalln(err)
		}
		prv, err := ucspi.PrivateKeyFromFile(*prvPath)
		if err != nil {
			log.Fatalln(err)
		}
		cfg.Certificates = []tls.Certificate{{
			Certificate: [][]byte{crtRaw},
			PrivateKey:  prv,
		}}
	}
	if *casPath != "" {
		var err error
		cfg.RootCAs, err = ucspi.CertPoolFromFile(*casPath)
		if err != nil {
			log.Fatalln(err)
		}
	}
	if *fpr != "" {
		hshOur, err := hex.DecodeString(*fpr)
		if err != nil {
			log.Fatalln(err)
		}
		cfg.VerifyPeerCertificate = func(rawCerts [][]byte, verifiedChains [][]*x509.Certificate) error {
			spki := verifiedChains[0][0].RawSubjectPublicKeyInfo
			hshTheir := sha256.Sum256(spki)
			if bytes.Compare(hshOur, hshTheir[:]) != 0 {
				return errors.New("server certificate's SPKI hash mismatch")
			}
			return nil
		}
	}

	conn := &ucspi.Conn{R: os.NewFile(6, "R"), W: os.NewFile(7, "W")}
	if conn.R == nil {
		log.Fatalln("no 6 file descriptor")
	}
	if conn.W == nil {
		log.Fatalln("no 7 file descriptor")
	}
	tlsConn := tls.Client(conn, cfg)
	if err := tlsConn.Handshake(); err != nil {
		log.Fatalln(err)
	}
	connState := tlsConn.ConnectionState()
	if *onlyShow {
		fmt.Fprintf(
			os.Stderr,
			"Version: %s\nCipherSuite: %s\n",
			ucspi.TLSVersion(connState.Version),
			tls.CipherSuiteName(connState.CipherSuite),
		)
		for _, cert := range connState.PeerCertificates {
			os.Stderr.WriteString("\n")
			pem.Encode(os.Stderr, &pem.Block{Type: "CERTIFICATE", Bytes: cert.Raw})
			hsh := sha256.Sum256(cert.RawSubjectPublicKeyInfo)
			fmt.Fprintf(
				os.Stderr, "Issuer: %s\nSubject: %s\nValidity: %s / %s\nSPKI SHA256: %s\n",
				cert.Issuer, cert.Subject, cert.NotBefore, cert.NotAfter,
				hex.EncodeToString(hsh[:]),
			)
		}
		return
	}
	dn := connState.PeerCertificates[0].Subject.String()

	args := flag.Args()
	cmd := exec.Command(args[0], args[1:]...)
	rr, rw, err := os.Pipe()
	if err != nil {
		log.Fatalln(err)
	}
	wr, ww, err := os.Pipe()
	if err != nil {
		log.Fatalln(err)
	}
	cmd.ExtraFiles = []*os.File{nil, nil, nil, rr, ww}
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	cmd.Env = append(os.Environ(), "PROTO=TLS")
	cmd.Env = append(cmd.Env, "TLSREMOTEDN="+dn)
	if err = cmd.Start(); err != nil {
		log.Fatalln(err)
	}
	copiers := make(chan struct{})
	go func() {
		io.Copy(rw, tlsConn)
		rw.Close()
		close(copiers)
	}()
	go func() {
		io.Copy(tlsConn, wr)
	}()
	_, err = cmd.Process.Wait()
	<-copiers
	if err != nil {
		log.Fatalln(err)
	}
}