// cmkeytool -- dealing with KEKS-encoded keypairs 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 (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"hash"
	"io"
	"log"
	"os"
	"sort"
	"strings"
	"time"

	"go.cypherpunks.su/keks"
	"go.cypherpunks.su/keks/cm"
	mceliece6960119x25519 "go.cypherpunks.su/keks/cm/enc/mceliece6960119-x25519"
	sntrup761x25519 "go.cypherpunks.su/keks/cm/enc/sntrup761-x25519"
	cmhash "go.cypherpunks.su/keks/cm/hash"
	"go.cypherpunks.su/keks/cm/sign"
	ed25519blake2b "go.cypherpunks.su/keks/cm/sign/ed25519-blake2b"
	"go.cypherpunks.su/keks/cm/sign/gost"
	"go.cypherpunks.su/keks/cm/sign/slhdsa"
)

const (
	FdPubR = 4
	FdPubW = 5
	FdPrvR = 8
	FdPrvW = 9
)

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

func main() {
	flag.Usage = usage
	ku := make(map[string]*struct{})
	sub := make(map[string]string)
	flag.Func(
		"ku",
		"Optional key usage, can be specified multiple times",
		func(v string) error {
			ku[v] = nil
			return nil
		},
	)
	flag.Func(
		"sub",
		"Part of subject, key=value, can be specified multiple times",
		func(v string) error {
			s := strings.SplitN(v, "=", 2)
			if len(s) != 2 {
				return errors.New("invalid key=value")
			}
			sub[s[0]] = s[1]
			return nil
		},
	)
	sinceRaw := flag.String("since", "",
		"Optional notBefore, \"2006-01-02 15:04:05\" format")
	lifetime := flag.Uint("lifetime", 365,
		"Lifetime of the certification, days")
	algo := flag.String("algo", ed25519blake2b.Ed25519BLAKE2b,
		"Public key algorithm")
	verify := flag.Bool("verify", false, "Verify provided -pub with -ca-pub")
	doList := flag.Bool("list", false, "List available algorithms")

	flag.Parse()
	log.SetFlags(log.Lshortfile)

	if *doList {
		algos := []string{
			ed25519blake2b.Ed25519BLAKE2b,
			gost.GOST3410256A,
			gost.GOST3410512C,
			sntrup761x25519.SNTRUP761X25519,
			mceliece6960119x25519.ClassicMcEliece6960119X25519,
			slhdsa.SLHDSASHAKE256s,
		}
		sort.Strings(algos)
		for _, s := range algos {
			fmt.Println(s)
		}
		return
	}

	fdPubR := os.NewFile(FdPubR, "pub-in")
	fdPubW := os.NewFile(FdPubW, "pub-out")
	fdPrvR := os.NewFile(FdPrvR, "prv-in")
	fdPrvW := os.NewFile(FdPrvW, "prv-out")

	var doCertify bool
	if len(sub) == 0 && !*verify {
		doCertify = true
	}

	var err error
	var since time.Time
	if *sinceRaw == "" {
		since = time.Now().UTC().Truncate(time.Second)
	} else {
		since, err = time.Parse(time.DateTime, *sinceRaw)
		if err != nil {
			log.Fatalln("while parsing -since:", err)
		}
	}
	till := since.Add(time.Duration(*lifetime) * 24 * time.Hour)

	var caPrv sign.Iface
	var caPubs []*sign.Signed
	if doCertify || *verify {
		data := mustReadAll(fdPubR)
		for len(data) > 0 {
			var signed *sign.Signed
			signed, data, err = sign.PubParse(data)
			if err != nil {
				log.Fatal(err)
			}
			caPubs = append(caPubs, signed)
		}
	}
	if doCertify {
		caPrv, _, err = sign.PrvParse(mustReadAll(fdPrvR))
		if err != nil {
			log.Fatal(err)
		}
	}

	var signed *sign.Signed
	if doCertify || *verify {
		signed, _, err = sign.PubParse(mustReadAll(os.Stdin))
		if err != nil {
			log.Fatal(err)
		}
	}

	if *verify {
		err = signed.CertificationVerify(caPubs, time.Now().UTC())
		if err != nil {
			log.Fatal(err)
		}
		return
	}

	var prvRaw []byte
	var pubData map[string]any
	if doCertify {
		pubData = (*signed.Data).(map[string]any)
	} else {
		var pub []byte
		switch *algo {
		case ed25519blake2b.Ed25519BLAKE2b:
			prvRaw, pub, err = ed25519blake2b.NewKeypair()
		case gost.GOST3410256A, gost.GOST3410512C:
			prvRaw, pub, err = gost.NewKeypair(*algo)
		case sntrup761x25519.SNTRUP761X25519:
			prvRaw, pub, err = sntrup761x25519.NewKeypair()
		case mceliece6960119x25519.ClassicMcEliece6960119X25519:
			prvRaw, pub, err = mceliece6960119x25519.NewKeypair()
		case slhdsa.SLHDSASHAKE256s:
			prvRaw, pub, err = slhdsa.NewKeypair(*algo)
		default:
			err = errors.New("unknown -algo specified")
		}
		if err != nil {
			log.Fatal(err)
		}
		{
			var buf bytes.Buffer
			if _, err = keks.Encode(&buf, sign.PrvMagic, nil); err != nil {
				log.Fatal(err)
			}
			if _, err = keks.Encode(&buf, cm.AV{A: *algo, V: prvRaw}, nil); err != nil {
				log.Fatal(err)
			}
			if _, err = io.Copy(fdPrvW, &buf); err != nil {
				log.Fatal(err)
			}
		}
		{
			pubData = map[string]any{
				"sub": sub,
				"pub": []cm.AV{{A: *algo, V: pub}},
			}
			var hasher hash.Hash
			switch *algo {
			case ed25519blake2b.Ed25519BLAKE2b, sntrup761x25519.SNTRUP761X25519:
				hasher = cmhash.ByName(cmhash.BLAKE2b256)
			case gost.GOST3410256A, gost.GOST3410512C:
				hasher = cmhash.ByName(cmhash.Streebog256)
			case mceliece6960119x25519.ClassicMcEliece6960119X25519, slhdsa.SLHDSASHAKE256s:
				hasher = cmhash.ByName(cmhash.SHAKE128)
			default:
				log.Fatal("unsupported algorithm")
			}
			_, err = keks.Encode(hasher, pubData["pub"], nil)
			if err != nil {
				log.Fatal(err)
			}
			pubData["id"] = hasher.Sum(nil)
			if err != nil {
				log.Fatal(err)
			}
		}
		if len(ku) > 0 {
			pubData["ku"] = ku
		}
	}

	{
		pubLoadAny := any(pubData)
		signed = &sign.Signed{TBS: sign.SDTBS{T: "pub"}, Data: &pubLoadAny}
	}

	if doCertify {
		if err = signed.CertifyWith(
			caPubs[0].PubData(), caPrv, since, till,
		); err != nil {
			log.Fatal(err)
		}
	}

	{
		var buf bytes.Buffer
		if _, err = keks.Encode(&buf, sign.PubMagic, nil); err != nil {
			log.Fatal(err)
		}
		if _, err = keks.Encode(&buf, signed, nil); err != nil {
			log.Fatal(err)
		}
		if _, err = io.Copy(fdPubW, &buf); err != nil {
			log.Fatal(err)
		}
	}
}
