package inspect

import (
	"crypto/rsa"
	"crypto/sha512"
	"crypto/x509"
	"encoding/binary"
	"errors"
	"strconv"
	"strings"
)

// PKCS1PublicKeyInfo attempts to parse a DER-form PKCS#1-encoded public RSA
// key, and returns information about it.
func PKCS1PublicKeyInfo(loc string, der []byte) Info {
	key, err := x509.ParsePKCS1PublicKey(der)
	if err != nil {
		return &BadInfo{
			Typ:        TypePublicKey,
			Loc:        loc,
			Underlying: err,
		}
	}

	return PublicKeyInfo(loc, key)
}

// PKCS1PublicKeyInfo attempts to parse a DER-form PKIX-encoded public RSA
// key, and returns information about it.
func PKIXPublicKeyInfo(loc string, der []byte) Info {
	key, err := x509.ParsePKIXPublicKey(der)
	if err != nil {
		return &BadInfo{
			Typ:        TypePublicKey,
			Loc:        loc,
			Underlying: err,
		}
	}

	rsakey, ok := key.(*rsa.PublicKey)
	if !ok {
		return &BadInfo{
			Typ:        TypePublicKey,
			Loc:        loc,
			Underlying: errors.New("PKIX public key type is not RSA"),
		}
	}

	return PublicKeyInfo(loc, rsakey)
}

// PublicKeyInfo returns structured information about the given RSA public key.
func PublicKeyInfo(loc string, key *rsa.PublicKey) *PublicKey {
	return &PublicKey{
		Loc:         loc,
		Bits:        key.N.BitLen(),
		Fingerprint: KeyFingerprint(key),
	}
}

// PublicKey holds structured information about an RSA public key. It
// implements Info.
type PublicKey struct {
	// Loc is the location the key was encountered.
	Loc string

	// Bits is the key size.
	Bits int

	// Fingerprint of the modulus, which can be used to compare keys.
	Fingerprint Fingerprint
}

// Type indicates this a a public key.
func (pub *PublicKey) Type() Type {
	return TypePublicKey
}

// Location returns the location data stored by PublicKeyInfo.
func (pub *PublicKey) Location() string {
	return pub.Loc
}

// Info returns structured information about the public key.
func (pub *PublicKey) Info() []Section {
	return []Section{
		pub.PublicKeyInfoSection(),
	}
}

// PublicKeyInfoSection returns the RSA public key specific information
// section.
func (pub *PublicKey) PublicKeyInfoSection() Section {
	return Section{
		Title: "RSA public key",
		Fields: []Field{
			Field{
				Key:   "Bits",
				Value: pub.Bits,
			},
			Field{
				Key:   "Fingerprint",
				Value: pub.Fingerprint,
			},
		},
	}
}

// Fingerprint of a key.
type Fingerprint []byte

// KeyFingerprint computes the fingerprint of an RSA key. It only requires the
// public section of the key for this.
func KeyFingerprint(key *rsa.PublicKey) Fingerprint {
	f := sha512.New512_224()
	var b [8]byte
	E := uint64(key.E)
	binary.BigEndian.PutUint64(b[:], E)
	f.Write(b[:])
	f.Write(key.N.Bytes())
	return Fingerprint(f.Sum(nil))
}

// String returns a nicely-formatted hex string that is easy to read/compare.
func (f Fingerprint) String() string {
	return FormatHexBytes([]byte(f))
}

// MarshalJSON returns a quoted hex string.
func (f Fingerprint) MarshalJSON() ([]byte, error) {
	return strconv.AppendQuote(nil, FormatHexBytes([]byte(f))), nil
}

// MarshalYAML returns a hex string.
func (f Fingerprint) MarshalYAML() (any, error) {
	return FormatHexBytes([]byte(f)), nil
}

const hexDig = "0123456789ABCDEF"

// FormatHexBytes returns a nicely-formatted hex string that is easy to read/compare.
func FormatHexBytes(b []byte) string {
	var s strings.Builder

	for len(b) > 1 {
		s.WriteByte(hexDig[b[0]>>4])
		s.WriteByte(hexDig[b[0]&15])
		s.WriteByte(':')
		b = b[1:]
	}
	if len(b) > 0 {
		s.WriteByte(hexDig[b[0]>>4])
		s.WriteByte(hexDig[b[0]&15])
	}
	return s.String()
}