//nolint:all // This is all test code
package common

import (
	"crypto"
	"crypto/ecdsa"
	"crypto/ed25519"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"fmt"
	"net"
)

// generateDeviceKeys generates the RSA key-pair a test device
// authenticates with.
func generateDeviceKeys() (*rsa.PrivateKey, *rsa.PublicKey, error) {
	privateKey, err := rsa.GenerateKey(rand.Reader, 1024)
	if err != nil {
		return nil, nil, err
	}
	return privateKey, &privateKey.PublicKey, nil
}

// signData signs data with SHA256+PKCS1v15 and returns the base64-encoded
// signature, the format deviceauth expects in the X-MEN-Signature header.
func signData(privateKey *rsa.PrivateKey, data []byte) (string, error) {
	hash := sha256.Sum256(data)
	signature, err := rsa.SignPKCS1v15(rand.Reader, privateKey, crypto.SHA256, hash[:])
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(signature), nil
}

// ExportPublicKeyPEM renders a public key (RSA, ECDSA or Ed25519) as the
// PKIX PEM string used in device auth requests. Panics if the key cannot be
// marshaled.
func ExportPublicKeyPEM(pubkey any) string {
	pubASN1, err := x509.MarshalPKIXPublicKey(pubkey)
	if err != nil {
		panic(fmt.Errorf("failed to marshal public key: %w", err))
	}
	pubBytes := pem.EncodeToMemory(&pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: pubASN1,
	})
	return string(pubBytes)
}

// EC curve identifiers mirroring the Python test suite's EC_CURVE_* constants
// (testutils/util/crypto.py).
var (
	ECCurveP224 = elliptic.P224()
	ECCurveP256 = elliptic.P256()
	ECCurveP384 = elliptic.P384()
	ECCurveP521 = elliptic.P521()
)

// generateECDeviceKeys generates an ECDSA key-pair on the given curve.
func generateECDeviceKeys(curve elliptic.Curve) (*ecdsa.PrivateKey, *ecdsa.PublicKey, error) {
	privateKey, err := ecdsa.GenerateKey(curve, rand.Reader)
	if err != nil {
		return nil, nil, err
	}
	return privateKey, &privateKey.PublicKey, nil
}

// generateEd25519DeviceKeys generates an Ed25519 key-pair.
func generateEd25519DeviceKeys() (ed25519.PrivateKey, ed25519.PublicKey, error) {
	pub, priv, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		return nil, nil, err
	}
	return priv, pub, nil
}

// signDataEC signs data with SHA256+ECDSA and returns the base64-encoded
// ASN.1/DER signature, the format deviceauth expects in the X-MEN-Signature
// header.
func signDataEC(privateKey *ecdsa.PrivateKey, data []byte) (string, error) {
	hash := sha256.Sum256(data)
	signature, err := ecdsa.SignASN1(rand.Reader, privateKey, hash[:])
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(signature), nil
}

// signDataEd25519 signs the raw (unhashed) data with Ed25519, per the Mender
// device auth signing convention, and returns the base64-encoded signature.
func signDataEd25519(privateKey ed25519.PrivateKey, data []byte) (string, error) {
	return base64.StdEncoding.EncodeToString(ed25519.Sign(privateKey, data)), nil
}

// Key kinds exercised by the device auth test suite, mirroring the key types
// generated by testutils/util/crypto.py.
const (
	KeyKindRSA     = "rsa"
	KeyKindECP224  = "ec-p224"
	KeyKindECP256  = "ec-p256"
	KeyKindECP384  = "ec-p384"
	KeyKindECP521  = "ec-p521"
	KeyKindEd25519 = "ed25519"
)

// KeyKinds lists all key kinds exercised by the device auth test suite.
var KeyKinds = []string{
	KeyKindRSA,
	KeyKindECP224,
	KeyKindECP256,
	KeyKindECP384,
	KeyKindECP521,
	KeyKindEd25519,
}

// KeyPair is a signing key-pair for a test device, abstracting over the key
// types deviceauth accepts (RSA, ECDSA P-224/256/384/521, Ed25519).
type KeyPair struct {
	Kind    string
	Private crypto.Signer
	Public  any
}

// NewKeyPair generates a key-pair of the given kind (one of the KeyKind*
// constants).
func NewKeyPair(kind string) (*KeyPair, error) {
	switch kind {
	case KeyKindRSA:
		priv, pub, err := generateDeviceKeys()
		if err != nil {
			return nil, err
		}
		return &KeyPair{Kind: kind, Private: priv, Public: pub}, nil
	case KeyKindECP224:
		return newECKeyPair(kind, ECCurveP224)
	case KeyKindECP256:
		return newECKeyPair(kind, ECCurveP256)
	case KeyKindECP384:
		return newECKeyPair(kind, ECCurveP384)
	case KeyKindECP521:
		return newECKeyPair(kind, ECCurveP521)
	case KeyKindEd25519:
		priv, pub, err := generateEd25519DeviceKeys()
		if err != nil {
			return nil, err
		}
		return &KeyPair{Kind: kind, Private: priv, Public: pub}, nil
	default:
		return nil, fmt.Errorf("unsupported key kind %q", kind)
	}
}

func newECKeyPair(kind string, curve elliptic.Curve) (*KeyPair, error) {
	priv, pub, err := generateECDeviceKeys(curve)
	if err != nil {
		return nil, err
	}
	return &KeyPair{Kind: kind, Private: priv, Public: pub}, nil
}

// PublicKeyPEM renders the public key as PKIX PEM.
func (k *KeyPair) PublicKeyPEM() string {
	return ExportPublicKeyPEM(k.Public)
}

// Sign signs data with the scheme appropriate for the key kind and returns
// the base64-encoded signature for the X-MEN-Signature header.
func (k *KeyPair) Sign(data []byte) (string, error) {
	return SignAuthRequest(k.Private, data)
}

// SignAuthRequest signs data with the scheme appropriate for the concrete
// type of privateKey (RSA, ECDSA or Ed25519) and returns the base64-encoded
// signature for the X-MEN-Signature header.
func SignAuthRequest(privateKey crypto.Signer, data []byte) (string, error) {
	switch priv := privateKey.(type) {
	case *rsa.PrivateKey:
		return signData(priv, data)
	case *ecdsa.PrivateKey:
		return signDataEC(priv, data)
	case ed25519.PrivateKey:
		return signDataEd25519(priv, data)
	default:
		return "", fmt.Errorf("unsupported private key type %T", priv)
	}
}

// RandomMAC returns a random unicast, locally administered MAC address for
// use as device identity data.
func RandomMAC() (net.HardwareAddr, error) {
	mac := make([]byte, 6)
	_, err := rand.Read(mac)
	if err != nil {
		return nil, err
	}

	mac[0] &= 0xfe // Set to Unicast
	mac[0] |= 0x02 // Set to Locally Administered

	return mac, nil
}
