// cmsigtool -- dealing with KEKS-encoded cm/signed utility
// Copyright (C) 2024-2025 Sergey Matveev <stargrave@stargrave.org>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser 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 Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this program.  If not, see <http://www.gnu.org/licenses/>.

package main

import (
	"bufio"
	"bytes"
	"encoding/hex"
	"flag"
	"hash"
	"io"
	"log"
	"os"
	"time"

	"go.cypherpunks.su/keks"
	cmballoon "go.cypherpunks.su/keks/cm/enc/balloon"
	cmhash "go.cypherpunks.su/keks/cm/hash"
	"go.cypherpunks.su/keks/cm/sign"
	"go.cypherpunks.su/keks/cm/sign/mode"
	"go.cypherpunks.su/keks/schema"
	"go.cypherpunks.su/keks/types"
)

const (
	FdPubR = 4
	FdPrvR = 8
)

const BlobChunkLen = 128 * 1024

func mustReadAll(r io.Reader) []byte {
	data, err := io.ReadAll(r)
	if err != nil {
		log.Fatal(err)
	}
	return data
}

func main() {
	log.SetFlags(log.Lshortfile)
	flag.Usage = usage
	typ := flag.String("t", "data", "Set/check the load type")
	verify := flag.Bool("v", false, "Do verification")
	var encryptedTo [][]byte
	flag.Func("encrypted-to", "Set/check encrypted-to, hex", func(v string) error {
		to, err := hex.DecodeString(v)
		if err != nil {
			return err
		}
		encryptedTo = append(encryptedTo, to)
		return nil
	})
	detached := flag.Bool("d", false, "Detached data mode")
	noWhen := flag.Bool("no-when", false, `Do not include "when"`)
	doMerkle := flag.Bool("merkle", false, "Use Merkle-tree based hasher")
	flag.Parse()

	fdPubR := os.NewFile(FdPubR, "pub-in")
	fdPrvR := os.NewFile(FdPrvR, "prv-in")

	var pubs []*sign.PubData
	var err error
	{
		data := mustReadAll(fdPubR)
		for len(data) > 0 {
			var signed *sign.Signed
			signed, data, err = sign.PubParse(data)
			if err != nil {
				log.Fatal(err)
			}
			pubs = append(pubs, signed.PubData())
		}
		fdPubR.Close()
	}

	stdin := bufio.NewReaderSize(os.Stdin, BlobChunkLen)
	if *verify {
		decoder := keks.NewDecoderFromReader(stdin, nil)
		var t types.Type
		t, err = decoder.Parse()
		if err != nil {
			log.Fatal(err)
		}
		if t != types.Magic || decoder.Iter().Magic() != sign.SignedMagic {
			log.Fatal("wrong magic")
		}
		decoder = keks.NewDecoderFromReader(stdin, nil)
		var v any
		v, err = decoder.Decode()
		if err != nil {
			log.Fatal(err)
		}
		var prehash cmhash.Prehash
		err = schema.Check("prehash", cmhash.PrehashSchemas, v)
		if err == nil {
			err = decoder.UnmarshalStruct(&prehash)
		}
		var signed sign.Signed
		hashers := make(map[string]hash.Hash)
		if err == nil && prehash.T == cmhash.PrehashT {
			dsts := make([]io.Writer, 0, len(prehash.Algos)+1)
			dsts = append(dsts, os.Stdout)
			for algo := range prehash.Algos {
				hasher := cmhash.ByName(algo)
				if hasher == nil {
					log.Fatalln("prehash: unsupported algorithm:", algo)
				}
				hashers[algo] = hasher
				dsts = append(dsts, hasher)
			}
			var blob *keks.BlobDecoder
			blob, err = keks.NewBlobDecoder(stdin, 1<<32)
			if err != nil {
				log.Fatal(err)
			}
			var chunk []byte
			mw := io.MultiWriter(dsts...)
			for {
				chunk, err = blob.Next()
				if err != nil {
					if err != io.EOF {
						log.Fatal(err)
					}
					break
				}
				if _, err = io.Copy(mw, bytes.NewReader(chunk)); err != nil {
					log.Fatal(err)
				}
			}
			decoder = keks.NewDecoderFromReader(stdin, nil)
			v, err = decoder.Decode()
			if err != nil {
				log.Fatal(err)
			}
		}
		err = schema.Check("signed", sign.SignedSchemas, v)
		if err != nil {
			log.Fatal(err)
		}
		err = decoder.UnmarshalStruct(&signed)
		if err != nil {
			log.Fatal(err)
		}
		if signed.TBS.T != *typ {
			log.Fatalln("differing load type:", signed.TBS.T)
		}
		for _, sig := range signed.Sigs {
			sid := sig.TBS["sid"].([]byte)
			var signerFound bool
			for _, pub := range pubs {
				if !bytes.Equal(sid, pub.Id) {
					continue
				}
				signerFound = true
				var sigTBS *sign.SigTBS
				sigTBS, err = sig.TBSGet()
				if err != nil {
					log.Fatal(err)
				}
				if len(encryptedTo) > 0 {
					if len(sigTBS.EncryptedTo) == 0 {
						log.Fatalln(hex.EncodeToString(sid), "missing encrypted-to")
					}
					found := false
					for _, their := range sigTBS.EncryptedTo {
						for _, our := range encryptedTo {
							if bytes.Equal(our, their) {
								found = true
								break
							}
						}
					}
					if !found {
						log.Fatalln(
							hex.EncodeToString(sid),
							"corresponding encrypted-to not found")
					}
				}
				var hasher hash.Hash
				if prehash.T == "" {
					hasher = cmhash.ByName(sig.Sign.A)
					if _, err = io.Copy(hasher, stdin); err != nil {
						log.Fatal(hex.EncodeToString(sid), err)
					}
				} else {
					var found bool
					hasher, found = hashers[sig.Sign.A]
					if !found {
						log.Fatalln(hex.EncodeToString(sid), "no hasher in prehash")
					}
				}
				if err = signed.CertificationCheckSignatureFrom(
					pub, &hasher,
				); err != nil {
					log.Fatalln(hex.EncodeToString(sid), err)
				}
				break
			}
			if !signerFound {
				log.Fatalln(hex.EncodeToString(sid), "can not find signer")
			}
		}
	} else {
		var signer sign.Iface
		{
			prvRaw := mustReadAll(fdPrvR)
			fdPrvR.Close()
			prvRaw, err = cmballoon.PossibleInteractiveDecrypt(prvRaw)
			if err != nil {
				log.Fatal(err)
			}
			signer, _, err = sign.PrvParse(prvRaw)
		}
		if err != nil {
			log.Fatal(err)
		}
		if *doMerkle {
			err = signer.SetMode(mode.Merkle)
		} else {
			err = signer.SetMode(mode.Prehash)
		}
		if err != nil {
			log.Fatal(err)
		}

		if _, err = keks.Encode(os.Stdout, sign.SignedMagic, nil); err != nil {
			log.Fatal(err)
		}
		if *detached {
			if _, err = io.Copy(*signer.Prehasher(), stdin); err != nil {
				log.Fatal(err)
			}
		} else {
			if _, err = keks.Encode(os.Stdout, cmhash.Prehash{
				T:     cmhash.PrehashT,
				Algos: map[string]*struct{}{signer.Algo(): nil},
			}, nil); err != nil {
				log.Fatal(err)
			}
			pr, pw := io.Pipe()
			br := keks.BlobReader{R: pr, ChunkLen: BlobChunkLen}
			copyErr := make(chan error)
			go func() {
				_, e := io.Copy(io.MultiWriter(pw, *signer.Prehasher()), stdin)
				pw.Close()
				copyErr <- e
			}()
			_, err = keks.Encode(os.Stdout, br, nil)
			if err != nil {
				log.Fatal(err)
			}
			if err = <-copyErr; err != nil {
				log.Fatal(err)
			}
		}
		var signed sign.Signed
		signed.TBS.T = *typ
		tbs := make(map[string]any)
		if !*noWhen {
			tbs["when"] = time.Now().UTC().Truncate(time.Millisecond)
		}
		if len(encryptedTo) > 0 {
			tbs["encrypted-to"] = encryptedTo
		}
		if err = signed.SignWith(pubs[0], signer, tbs); err != nil {
			log.Fatal(err)
		}
		if _, err = keks.Encode(os.Stdout, signed, nil); err != nil {
			log.Fatal(err)
		}
	}
}
