package main

import (
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"github.com/pkg/errors"
	"gopkg.in/square/go-jose.v2"
	"io/ioutil"
)

// TODO: CRL validation

type KeyValidator interface {
	Validate(jose.JSONWebKey) error
	LoadRootPEM(string) error
}

type keyValidator struct {
	pkiSubject string
	algorithms *stringSet
	roots      *x509.CertPool
}

func NewKeyValidator(subject string) KeyValidator {
	return &keyValidator{
		pkiSubject: subject,
		algorithms: NewStringSet("PS256", "PS385", "PS512"),
		roots:      x509.NewCertPool(),
	}
}

func (v *keyValidator) LoadRootPEM(filename string) error {
	pem_data, err := ioutil.ReadFile(filename)
	if err != nil {
		return errors.WithStack(err)
	}

	pem_block, _ := pem.Decode(pem_data)
	if pem_block == nil {
		return errors.Errorf("PEM decode failed")
	}

	cert, err := x509.ParseCertificate(pem_block.Bytes)
	if err != nil {
		return errors.WithStack(err)
	}

	v.roots.AddCert(cert)

	return nil
}

func (v *keyValidator) Validate(key jose.JSONWebKey) error {
	pk, ok := key.Key.(*rsa.PublicKey)
	if !ok {
		return errors.Errorf("Key type is not RSA")
	}

	if !v.algorithms.Contains(key.Algorithm) {
		return errors.Errorf("Key algorithm is not supported")
	}

	cert := key.Certificates[0]
	cpk, ok := cert.PublicKey.(*rsa.PublicKey)
	if !ok {
		return errors.Errorf("Public key is not RSA")
	}

	if cpk.N.BitLen() < 2048 {
		return errors.Errorf("Key length less than 2048 bits")
	}

	if cert.KeyUsage&x509.KeyUsageDigitalSignature != 1 {
		return errors.Errorf("Certificate not valid for digital signatures")
	}

	err := v.validateCertificateChain(key.Certificates)
	if err != nil {
		return errors.WithStack(err)
	}

	err = v.validateCertificateCRL(cert)
	if err != nil {
		return errors.WithStack(err)
	}

	err = v.validatePublicKeyInCertificate(pk, cpk)
	if err != nil {
		return errors.WithStack(err)
	}

	return nil
}

// TODO
// Fetch CRL from distrubtion point in cert
// Validate CRL signed by trusted CA
// Validate cert not in CRL
func (v *keyValidator) validateCertificateCRL(cert *x509.Certificate) error {
	return nil
}

func (v *keyValidator) validateCertificateChain(chain []*x509.Certificate) error {
	vo := x509.VerifyOptions{
		Roots:     v.roots,
		KeyUsages: []x509.ExtKeyUsage{x509.ExtKeyUsageAny},
	}

	if len(chain) > 1 {
		ip := x509.NewCertPool()
		for _, i := range chain[1:] {
			ip.AddCert(i)
		}

		vo.Intermediates = ip
	}

	chains, err := chain[0].Verify(vo)
	if err != nil {
		return errors.WithStack(err)
	}

	if len(chains) <= 0 {
		return errors.Errorf("No valid certificate chains found")
	}

	if chain[0].Subject.CommonName != v.pkiSubject {
		return errors.Errorf("Invalid certificate subject name")
	}

	return nil
}

// validate first item of x5c matches n and e
func (v *keyValidator) validatePublicKeyInCertificate(pk *rsa.PublicKey, cpk *rsa.PublicKey) error {
	if cpk.E != pk.E {
		return errors.Errorf("E in key and E in cert do not match")
	}

	if pk.N.Cmp(cpk.N) != 0 {
		return errors.Errorf("N in key and N in cert do not match")
	}

	return nil
}