package sshutil

import (
	"crypto/sha256"
	"encoding/base64"
	"math/big"

	"golang.org/x/crypto/ssh"
)

func KeyBits(key ssh.PublicKey) int {
	if key == nil {
		return 0
	}
	switch key.Type() {
	case ssh.KeyAlgoECDSA256:
		return 256
	case ssh.KeyAlgoSKECDSA256:
		return 256
	case ssh.KeyAlgoECDSA384:
		return 384
	case ssh.KeyAlgoECDSA521:
		return 521
	case ssh.KeyAlgoED25519:
		return 256
	case ssh.KeyAlgoSKED25519:
		return 256
	case ssh.KeyAlgoRSA:
		var w struct {
			Name string
			E    *big.Int
			N    *big.Int
			Rest []byte `ssh:"rest"`
		}
		_ = ssh.Unmarshal(key.Marshal(), &w)
		return w.N.BitLen()
	default:
		return 0
	}
}

func KeyType(key ssh.PublicKey) string {
	if key == nil {
		return "<nil>"
	}
	switch key.Type() {
	case ssh.KeyAlgoECDSA256, ssh.KeyAlgoSKECDSA256, ssh.KeyAlgoECDSA384, ssh.KeyAlgoECDSA521:
		return "ECDSA"
	case ssh.KeyAlgoED25519, ssh.KeyAlgoSKED25519:
		return "ED25519"
	case ssh.KeyAlgoRSA:
		return "RSA"
	default:
		return key.Type()
	}
}

func KeyFingerprint(key ssh.PublicKey) string {
	h := sha256.New()
	h.Write(key.Marshal())
	return "SHA256:" + base64.RawStdEncoding.EncodeToString(h.Sum(nil))
}