jx/internal/resource/pki.go

// Copyright 2024 Matthew Rich <matthewrich.conf@gmail.com>. All rights reserved.

package resource

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"gopkg.in/yaml.v3"
	"net/url"
	"path/filepath"
	"gitea.rosskeen.house/rosskeen.house/machine"
	"decl/internal/codec"
	"decl/internal/ext"
	"decl/internal/transport"
	"decl/internal/data"
	"decl/internal/folio"
	"crypto/x509"
	"crypto/x509/pkix"
	"crypto/rsa"
	"crypto/rand"
	"encoding/pem"
	"encoding/json"
	"time"
	"math/big"
	"io"
	"strings"
)

const (
	PKITypeName TypeName = "pki"
)

// Describes the type of certificate file the resource represents
type EncodingType string

// Supported file types
const (
	EncodingTypePem EncodingType = "pem"
)

var ErrPKIInvalidEncodingType error = errors.New("Invalid EncodingType")
var ErrPKIFailedDecodingPemBlock error = errors.New("Failed decoding pem block")

func init() {
	ResourceTypes.Register([]string{"pki"}, func(u *url.URL) (data.Resource) {
		k := NewPKI()
		if u != nil {
			if err := folio.CastParsedURI(u).ConstructResource(k); err != nil {
                                panic(err)
                        }
			ref := folio.ResourceReference(filepath.Join(u.Hostname(), u.Path))
			if len(ref) > 0 {
				k.PrivateKeyRef = ref
			}
		}
		return k
	})
}

type Certificate struct {
	*x509.Certificate `yaml:",inline" json:",inline"`
	Raw                     []byte `yaml:"-" json:"-"`
	RawTBSCertificate       []byte `yaml:"-" json:"-"` 
	RawSubjectPublicKeyInfo []byte `yaml:"-" json:"-"` 
	RawSubject              []byte `yaml:"-" json:"-"` 
	RawIssuer               []byte `yaml:"-" json:"-"` 
}

type PKI struct {
	*Common `json:",inline" yaml:",inline"`
	stater machine.Stater `json:"-" yaml:"-"`
	PrivateKeyPem string `json:"privatekey,omitempty" yaml:"privatekey,omitempty"`
	PublicKeyPem string `json:"publickey,omitempty" yaml:"publickey,omitempty"`
	CertificatePem string `json:"certificate,omitempty" yaml:"certificate,omitempty"`

	PrivateKeyRef folio.ResourceReference `json:"privatekeyref,omitempty" yaml:"privatekeyref,omitempty"` // Describes a resource URI to read/write the private key content
	PublicKeyRef folio.ResourceReference `json:"publickeyref,omitempty" yaml:"publickeyref,omitempty"`  // Describes a resource URI to read/write the public key content 
	CertificateRef folio.ResourceReference `json:"certificateref,omitempty" yaml:"certificateref,omitempty"` // Describes a resource URI to read/write the certificate content

	SignedByRef folio.ResourceReference `json:"signedbyref,omitempty" yaml:"signedbyref,omitempty"` // Describes a resource URI for the signing cert

	privateKey *rsa.PrivateKey `json:"-" yaml:"-"`
	publicKey *rsa.PublicKey `json:"-" yaml:"-"`
	Values *Certificate `json:"values,omitempty" yaml:"values,omitempty"`
	Certificate *x509.Certificate `json:"-" yaml:"-"`

	Bits int `json:"bits" yaml:"bits"`
	EncodingType EncodingType `json:"type" yaml:"type"`
	Resources data.ResourceMapper `json:"-" yaml:"-"`
}

func NewPKI() *PKI {
	p := &PKI{ EncodingType: EncodingTypePem, Bits: 2048, Common: NewCommon(PKITypeName, false) }
	return p
}

func (k *PKI) SetResourceMapper(resources data.ResourceMapper) {
	slog.Info("PKI.SetResourceMapper()", "resources", resources)
	k.Resources = resources
}

func (k *PKI) Clone() data.Resource {
	return &PKI {
		Common: k.Common.Clone(),
		EncodingType: k.EncodingType,
	}
}

func (k *PKI) StateMachine() machine.Stater {
	if k.stater == nil {
		k.stater = StorageMachine(k)
	}
	return k.stater
}

func (k *PKI) Notify(m *machine.EventMessage) {
	ctx := context.Background()
	slog.Info("Notify()", k.Type(), k, "m", m)
	switch m.On {
	case machine.ENTERSTATEEVENT:
		switch m.Dest {
			case "start_stat":
				if statErr := k.ReadStat(); statErr == nil {
					if triggerErr := k.StateMachine().Trigger("exists"); triggerErr == nil {
						return
					}
				} else {
					if triggerErr := k.StateMachine().Trigger("notexists"); triggerErr == nil {
						return
					}
				}
			case "start_read":
				if _,readErr := k.Read(ctx); readErr == nil {
					if triggerErr := k.StateMachine().Trigger("state_read"); triggerErr == nil {
						return
					} else {
						//k.State = "absent"
						panic(triggerErr)
					}
			} else {
				//k.State = "absent"
				panic(readErr)
			}
		case "start_create":
			if ! (k.PrivateKeyRef.Exists() || k.PublicKeyRef.Exists() || k.CertificateRef.Exists()) {
				if e := k.Create(ctx); e == nil {
					if triggerErr := k.StateMachine().Trigger("created"); triggerErr == nil {
						return
					}
				}
			}
			//k.State = "absent"
		case "start_update":
			if e := k.Update(ctx); e == nil {
				if triggerErr := k.StateMachine().Trigger("updated"); triggerErr == nil {
					return
				}
			}
		case "start_delete":
			if deleteErr := k.Delete(ctx); deleteErr == nil {
				if triggerErr := k.StateMachine().Trigger("deleted"); triggerErr == nil {
					return
				} else {
					//k.State = "present"
					panic(triggerErr)
				}
			} else {
				//k.State = "present"
				panic(deleteErr)
			}
		case "absent":
			//k.State = "absent"
		case "present", "updated", "created", "read":
			//k.State = "present"
		}
	case machine.EXITSTATEEVENT:
	}
}

func (k *PKI) URI() string {
	u := k.PrivateKeyRef.Parse()
	if u.Scheme == "file" || u.Scheme == "pki" {
		return fmt.Sprintf("pki://%s", filepath.Join(u.Hostname(), u.Path))
	}
	return fmt.Sprintf("pki://%s", filepath.Join(u.Hostname(), u.RequestURI()))
}

func (k *PKI) Validate() error {
	return fmt.Errorf("failed")
}

func (k *PKI) Apply() error {
/*
	ctx := context.Background()
	switch k.State {
	case "absent":
		return k.Delete(ctx)
	case "present":
		return k.Create(ctx)
	}
*/

	return nil
}

func (k *PKI) Load(docData []byte, f codec.Format) (err error) {
	err = f.StringDecoder(string(docData)).Decode(k)
	return
}

func (k *PKI) LoadReader(r io.ReadCloser, f codec.Format) (err error) {
	err = f.Decoder(r).Decode(k)
	return
}

func (k *PKI) LoadString(docData string, f codec.Format) (err error) {
	err = f.StringDecoder(docData).Decode(k)
	return
}

func (k *PKI) LoadDecl(yamlResourceDeclaration string) error {
	return k.LoadString(yamlResourceDeclaration, codec.FormatYaml)
}

func (k *PKI) ReadStat() (err error) {
	var resourcesErr []string
	if ! k.PrivateKeyRef.Exists() {
		resourcesErr = append(resourcesErr, string(k.PrivateKeyRef))
	}

	if ! k.PublicKeyRef.Exists() {
		resourcesErr = append(resourcesErr, string(k.PublicKeyRef))
	}

	if ! k.CertificateRef.Exists() {
		resourcesErr = append(resourcesErr, string(k.CertificateRef))
	}

	if len(resourcesErr) > 0 {
		err = fmt.Errorf("PKI resources missing: %s", strings.Join(resourcesErr, ","))
//	k.State = "absent"
	}
	return
}

func (k *PKI) ResolveId(ctx context.Context) string {
	return string(k.PrivateKeyRef)
}

func (k *PKI) GenerateKey() (err error) {
	k.privateKey, err = rsa.GenerateKey(rand.Reader, k.Bits)
	return
}

func (k *PKI) PublicKey() {
	if k.privateKey != nil {
		k.publicKey = k.privateKey.Public().(*rsa.PublicKey)
	}
}

func (k *PKI) Encode() {
	var privFileStream, pubFileStream io.WriteCloser
	switch k.EncodingType {
	case EncodingTypePem:
		if len(k.PrivateKeyRef) > 0 {
			privFileStream, _ = k.PrivateKeyRef.Lookup(k.Resources).ContentWriterStream()
		} else {
			var privateKeyContent strings.Builder
			privFileStream = ext.WriteNopCloser(&privateKeyContent)
			defer func() { k.PrivateKeyPem = privateKeyContent.String() }()
		}
		defer privFileStream.Close()
		privEncodeErr := pem.Encode(privFileStream, &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(k.privateKey)})
		if privEncodeErr != nil {
			panic(privEncodeErr)
		}
		if len(k.PublicKeyRef) > 0 {
			pubFileStream, _ = k.PublicKeyRef.Lookup(k.Resources).ContentWriterStream()
		} else {
			var publicKeyContent strings.Builder
			pubFileStream = ext.WriteNopCloser(&publicKeyContent)
			defer func() { k.PublicKeyPem = publicKeyContent.String() }()
		}
		defer pubFileStream.Close()
		pubEncodeErr := pem.Encode(pubFileStream, &pem.Block{Type: "RSA PUBLIC KEY", Bytes: x509.MarshalPKCS1PublicKey(k.publicKey)})
		if pubEncodeErr != nil {
			panic(pubEncodeErr)
		}
	}
}

func (k *PKI) Decode() {
	slog.Info("PKI.Decode()", "privatekey", k.PrivateKeyRef, "publickey", k.PublicKeyRef, "certificate", k.CertificateRef)
	var err error
	switch k.EncodingType {
	case EncodingTypePem:
		if len(k.PrivateKeyRef) > 0 && k.PrivateKeyRef.Exists() {
			privReader := k.PrivateKeyRef.Lookup(k.Resources)
			privFileStream, _ := privReader.ContentReaderStream()
			defer privFileStream.Close()
			PrivateKeyPemData, readErr := io.ReadAll(privFileStream)
			if readErr != nil {
				panic(readErr)
			}
			k.PrivateKeyPem = string(PrivateKeyPemData)
			block, _ := pem.Decode(PrivateKeyPemData)
			if block != nil {
				k.privateKey, err = x509.ParsePKCS1PrivateKey(block.Bytes)
				if err != nil {
					panic(err)
				}
			} else {
				panic(ErrPKIFailedDecodingPemBlock)
			}
		}
		slog.Info("PKI.Decode() decoded private key", "error", err)
		if len(k.PublicKeyRef) > 0 && k.PublicKeyRef.Exists() {
			pubReader := k.PublicKeyRef.Lookup(k.Resources)
			pubFileStream, _ := pubReader.ContentReaderStream()
			defer pubFileStream.Close()
			PublicKeyPemData, readErr := io.ReadAll(pubFileStream)
			if readErr != nil {
				panic(err)
			}
			k.PublicKeyPem = string(PublicKeyPemData)
			block, _ := pem.Decode(PublicKeyPemData)
			if block != nil {
				k.publicKey, err = x509.ParsePKCS1PublicKey(block.Bytes)
				if err != nil {
					panic(err)
				}
			} else {
				panic(ErrPKIFailedDecodingPemBlock)
			}
		}
		slog.Info("PKI.Decode() decoded public key", "publickey", k.PublicKeyPem, "error", err)
		if len(k.CertificateRef) > 0 && k.CertificateRef.Exists() {
			certReader := k.CertificateRef.Lookup(k.Resources)
			certFileStream, _ := certReader.ContentReaderStream()
			if certFileStream != nil {
				defer certFileStream.Close()
				CertificatePemData, readErr := io.ReadAll(certFileStream)
				if readErr != nil {
					panic(readErr)
				}
				slog.Info("PKI.Decode() certificate", "pem", CertificatePemData, "error", err)
				k.CertificatePem = string(CertificatePemData)
				block, _ := pem.Decode(CertificatePemData)
				if block != nil {
					k.Certificate, err = x509.ParseCertificate(block.Bytes)
					if err != nil {
						panic(err)
					}
				} else {
					panic(ErrPKIFailedDecodingPemBlock)
				}
			}
		}
	}
	slog.Info("PKI.Decode()", "error", err)
}

func (k *PKI) ContentReaderStream() (*transport.Reader, error) {
	return nil, nil
}

func (k *PKI) ContentWriterStream() (*transport.Writer, error) {
	return nil, nil
}

func (k *PKI) SignedBy() (cert *x509.Certificate, pub *rsa.PublicKey, priv *rsa.PrivateKey) {
	if len(k.SignedByRef) > 0 {
		r := k.SignedByRef.Dereference(k.Resources)
		if r != nil {
			pkiResource := r.(*PKI)
			return pkiResource.Certificate, pkiResource.publicKey, pkiResource.privateKey
		}
	}
	return nil, nil, nil
}

func (k *PKI) CertConfig() (certTemplate *x509.Certificate) {
	if k.Values != nil {
		certTemplate = k.Values.Certificate
	} else {
		var caEpoch int64 = 1721005200000
		certTemplate = &x509.Certificate {
			SerialNumber: big.NewInt(time.Now().UnixMilli() - caEpoch),
			Subject: pkix.Name {
				Organization: []string{""},
				Country: []string{""},
				Province: []string{""},
				Locality: []string{""},
				StreetAddress: []string{""},
				PostalCode: []string{""},
			},
			NotBefore: time.Now(),
			NotAfter: time.Now().AddDate(10, 0, 0),
			KeyUsage: x509.KeyUsageDigitalSignature,
			ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		}
	}
	if k.config != nil {
		if value, configErr := k.config.GetValue("certtemplate"); configErr == nil {
			certTemplate = value.(*x509.Certificate)
		}
	}
	slog.Info("PKI.CertConfig()", "template", certTemplate)
	return
}

func (k *PKI) GenerateCertificate() (error) {
	var signingCert *x509.Certificate
	var signingPubKey *rsa.PublicKey
	var signingPrivKey *rsa.PrivateKey
	var certFileStream io.WriteCloser

	certTemplate := k.CertConfig()
	signingCert, signingPubKey, signingPrivKey = k.SignedBy()
	if signingCert != nil && signingPubKey != nil && signingPrivKey != nil {
		slog.Info("PKI.Certificate()", "signedby", k.SignedByRef)
	} else {
		signingCert = certTemplate
		signingPubKey = k.publicKey
		signingPrivKey = k.privateKey
	}
	if k.Certificate != nil {
		certTemplate = k.Certificate
	}
 	cert, err := x509.CreateCertificate(rand.Reader, certTemplate, signingCert, signingPubKey, signingPrivKey)
	if err != nil {
		slog.Error("PKI.Certificate() - x509.CreateCertificate", "cert", cert, "error", err)
		return err
	}

	if len(k.CertificateRef) > 0 {
		certFileStream, _ = k.CertificateRef.Lookup(k.Resources).ContentWriterStream()
	} else {
			var certContent strings.Builder
			certFileStream = ext.WriteNopCloser(&certContent)
			defer func() { k.CertificatePem = certContent.String() }()
	}
	defer certFileStream.Close()
	certEncodeErr := pem.Encode(certFileStream, &pem.Block{Type: "CERTIFICATE", Bytes: cert})
	return certEncodeErr
}

func (k *PKI) Create(ctx context.Context) (err error) {
	if err = k.GenerateKey(); err == nil {
		k.PublicKey()
		k.Encode()
		if err = k.GenerateCertificate(); err == nil {
			return
		}
	}
	return
}

func (f *PKI) Delete(ctx context.Context) error {
	return nil
}

func (k *PKI) Read(ctx context.Context) ([]byte, error) {
	k.Decode()
	return yaml.Marshal(k)
}

func (k *PKI) Update(ctx context.Context) (err error) {
	if err = k.GenerateKey(); err == nil {
		k.PublicKey()
		k.Encode()
		if err = k.GenerateCertificate(); err == nil {
			return
		}
	}
	return
}

func (k *PKI) Type() string { return "pki" }

func (t *EncodingType) UnmarshalValue(value string) error {
	switch value {
	case string(EncodingTypePem):
		*t = EncodingType(value)
		return nil
	default:
		return ErrPKIInvalidEncodingType
	}
}

func (t *EncodingType) UnmarshalJSON(data []byte) error {
	var s string
	if unmarshalEncodingTypeErr := json.Unmarshal(data, &s); unmarshalEncodingTypeErr != nil {
		return unmarshalEncodingTypeErr
	}
	return t.UnmarshalValue(s)
}

func (t *EncodingType) UnmarshalYAML(value *yaml.Node) error {
	var s string
	if err := value.Decode(&s); err != nil {
		return err
	}
	return t.UnmarshalValue(s)
}