package main

import (
	"crypto"
	"crypto/ecdsa"
	"crypto/rand"
	"crypto/sha256"
	"crypto/x509"
	"encoding/asn1"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"log"
	"math/big"
	"os"

	"github.com/google/go-tpm/legacy/tpm2"
	"github.com/google/go-tpm/tpmutil"
)

func generateSrk(f *os.File) {
	tmpl := tpm2.Public{
		Type:    tpm2.AlgRSA,
		NameAlg: tpm2.AlgSHA256,
		Attributes: tpm2.FlagFixedTPM | // Key can't leave the TPM.
			tpm2.FlagFixedParent | // Key can't change parent.
			tpm2.FlagSensitiveDataOrigin | // Key created by the TPM (not imported).
			tpm2.FlagUserWithAuth | // Uses (empty) password.
			tpm2.FlagNoDA | // This flag doesn't do anything, but it's in the spec.
			tpm2.FlagRestricted | // Key used for TPM challenges, not general decryption.
			tpm2.FlagDecrypt, // Key can be used to decrypt data.
		RSAParameters: &tpm2.RSAParams{
			Symmetric:  &tpm2.SymScheme{Alg: tpm2.AlgAES, KeyBits: 128, Mode: tpm2.AlgCFB},
			KeyBits:    2048,
			ModulusRaw: make([]byte, 256),
		},
	}

	srk, _, err := tpm2.CreatePrimary(f, tpm2.HandleOwner, tpm2.PCRSelection{}, "", "", tmpl)
	if err != nil {
		log.Fatalf("creating srk: %v", err)
	}
	out, err := tpm2.ContextSave(f, srk)
	if err != nil {
		log.Fatalf("saving context: %v", err)
	}
	if err := os.WriteFile("srk.ctx", out, 0644); err != nil {
		log.Fatalf("writing context: %v", err)
	}
}

func generateAppKey(f *os.File) {
	srkCtx, err := os.ReadFile("srk.ctx")
	if err != nil {
		log.Fatalf("read srk: %v", err)
	}
	srk, err := tpm2.ContextLoad(f, srkCtx)
	if err != nil {
		log.Fatalf("load srk: %v", err)
	}

	tmpl := tpm2.Public{
		Type:    tpm2.AlgECC,
		NameAlg: tpm2.AlgSHA256,
		Attributes: tpm2.FlagFixedTPM | // Key can't leave the TPM.
			tpm2.FlagFixedParent | // Key can't change parent.
			tpm2.FlagSensitiveDataOrigin | // Key created by the TPM (not imported).
			tpm2.FlagUserWithAuth | // Uses (empty) password.
			tpm2.FlagSign, // Key can be used to sign data.
		ECCParameters: &tpm2.ECCParams{
			Sign:    &tpm2.SigScheme{Alg: tpm2.AlgECDSA, Hash: tpm2.AlgSHA256},
			CurveID: tpm2.CurveNISTP256,
			Point: tpm2.ECPoint{
				XRaw: make([]byte, 32),
				YRaw: make([]byte, 32),
			},
		},
	}

	privBlob, pubBlob, _, hash, ticket, err := tpm2.CreateKey(f, srk, tpm2.PCRSelection{}, "", "", tmpl)
	if err != nil {
		log.Fatalf("create aik: %v", err)
	}
	log.Printf("privBlob: %x\n", privBlob)
	log.Printf("pubBlob: %x\n", pubBlob)
	log.Printf("hash: %x\n", hash)
	log.Printf("ticket: %v %s\n", ticket.Type, ticket.Digest)

	appKey, _, err := tpm2.Load(f, srk, "", pubBlob, privBlob)
	if err != nil {
		log.Fatalf("load app key: %v", err)
	}

	// Write key context to disk.
	appKeyCtx, err := tpm2.ContextSave(f, appKey)
	if err != nil {
		log.Fatalf("saving context: %v", err)
	}
	if err := os.WriteFile("app.ctx", appKeyCtx, 0644); err != nil {
		log.Fatalf("writing context: %v", err)
	}
}

func fileExists(filename string) bool {
	_, err := os.Stat(filename)
	if err == nil {
		return true
	}
	if errors.Is(err, os.ErrNotExist) {
		return false
	}
	panic("Should not happen")
}

type tpmSinger struct {
	tpm io.ReadWriter
	h   tpmutil.Handle
	pub crypto.PublicKey
}

func (s *tpmSinger) Public() crypto.PublicKey {
	return s.pub
}

func (s *tpmSinger) Sign(r io.Reader, digest []byte, opts crypto.SignerOpts) ([]byte, error) {
	sig, err := tpm2.Sign(s.tpm, s.h, "", digest, nil, nil)
	if err != nil {
		return nil, fmt.Errorf("signing data: %v", err)
	}
	if sig.RSA != nil {
		return sig.RSA.Signature, nil
	}
	if sig.ECC != nil {
		return asn1.Marshal(struct {
			R *big.Int
			S *big.Int
		}{sig.ECC.R, sig.ECC.S})
	}
	return nil, fmt.Errorf("unsupported signature type: %v", sig.Alg)
}

func newSigner(tpm io.ReadWriter, h tpmutil.Handle) (*tpmSinger, error) {
	tpmPub, _, _, err := tpm2.ReadPublic(tpm, h)
	if err != nil {
		return nil, fmt.Errorf("read public blob: %v", err)
	}
	pub, err := tpmPub.Key()
	if err != nil {
		return nil, fmt.Errorf("decode public key: %v", err)
	}
	return &tpmSinger{tpm, h, pub}, nil
}

func getSinger(f *os.File) (*tpmSinger, error) {
	keyCtx, err := os.ReadFile("app.ctx")
	if err != nil {
		log.Fatalf("read app key: %v", err)
	}
	key, err := tpm2.ContextLoad(f, keyCtx)
	if err != nil {
		log.Fatalf("load app key: %v", err)
	}

	return newSigner(f, key)
}

// reference: https://ericchiang.github.io/post/tpm-keys/
func main() {
	f, err := os.OpenFile("/dev/tpmrm0", os.O_RDWR, 0)
	if err != nil {
		log.Fatalf("opening tpm: %v", err)
	}
	defer f.Close()

	if !fileExists("srk.ctx") {
		generateSrk(f)
	}
	if !fileExists("app.ctx") {
		generateAppKey(f)
	}

	priv, err := getSinger(f)
	if err != nil {
		log.Fatalf("create signer: %v", err)
	}

	publicKeyDer, err := x509.MarshalPKIXPublicKey(priv.Public())
	if err != nil {
		log.Fatalf("public key: %v", err)
	} else {
		log.Printf("public key pem: %s\n", base64.StdEncoding.EncodeToString(publicKeyDer))
	}

	msg := []byte("Hello world!")
	digest := sha256.Sum256(msg)

	sig, err := priv.Sign(rand.Reader, digest[:], crypto.SHA256)
	if err != nil {
		log.Fatalf("signing data: %v", err)
	}
	log.Printf("sig: %s\n", hex.EncodeToString(sig))
	pub, ok := priv.Public().(*ecdsa.PublicKey)
	if !ok {
		log.Fatalf("expected ecdsa.PublicKey got: %T", priv.Public())
	}
	if !ecdsa.VerifyASN1(pub, digest[:], sig) {
		log.Fatalf("failed to verify signature")
	}
}