package main

// Keys, by their public half only. A fingerprint is computed here from the public key's bytes, as
// ssh-keygen -l does (SHA256 of the key blob, unpadded base64); a private key is never read by this
// process — its first line, the PEM header, says it is one, and ssh-keygen is asked about it.

import (
	"crypto/sha256"
	"encoding/base64"
	"encoding/binary"
	"fmt"
	"math/big"
	"strings"
)

// KeyTypes are the public key algorithms an authorized_keys or known_hosts line can carry.
var KeyTypes = map[string]bool{
	"ssh-ed25519": true, "ssh-rsa": true, "ssh-dss": true,
	"ecdsa-sha2-nistp256": true, "ecdsa-sha2-nistp384": true, "ecdsa-sha2-nistp521": true,
	"sk-ssh-ed25519@openssh.com": true, "sk-ecdsa-sha2-nistp256@openssh.com": true,
}

// PublicKey is one public key as a line carries it.
type PublicKey struct {
	Type        string `json:"type"`
	Bits        int    `json:"bits"`
	Fingerprint string `json:"fingerprint"`
	Comment     string `json:"comment"`
	Options     string `json:"options,omitempty"`
	Weak        string `json:"weak,omitempty"`
}

// ParseKeyLine reads one authorized_keys line (options, type, key, comment) or a public key file's.
func ParseKeyLine(line string) (PublicKey, error) {
	fields := fieldsQuoted(strings.TrimSpace(line))
	for i, f := range fields {
		if !KeyTypes[f] || i+1 >= len(fields) {
			continue
		}
		k := PublicKey{Type: f, Options: strings.Join(fields[:i], " "), Comment: strings.Join(fields[i+2:], " ")}
		blob, err := base64.StdEncoding.DecodeString(fields[i+1])
		if err != nil {
			return k, fmt.Errorf("the key after %s is not base64", f)
		}
		sum := sha256.Sum256(blob)
		k.Fingerprint = "SHA256:" + base64.RawStdEncoding.EncodeToString(sum[:])
		k.Bits = bitsOf(f, blob)
		switch {
		case f == "ssh-dss":
			k.Weak = "DSA: refused by current OpenSSH"
		case f == "ssh-rsa" && k.Bits > 0 && k.Bits < 3072:
			k.Weak = fmt.Sprintf("RSA of %d bits: below 3072", k.Bits)
		}
		return k, nil
	}
	return PublicKey{}, fmt.Errorf("no public key on this line")
}

// fieldsQuoted splits on spaces outside double quotes, as sshd reads an options field.
func fieldsQuoted(s string) []string {
	var out []string
	var cur strings.Builder
	quoted := false
	for _, ch := range s {
		switch {
		case ch == '"':
			quoted = !quoted
			cur.WriteRune(ch)
		case (ch == ' ' || ch == '\t') && !quoted:
			if cur.Len() > 0 {
				out = append(out, cur.String())
				cur.Reset()
			}
		default:
			cur.WriteRune(ch)
		}
	}
	if cur.Len() > 0 {
		out = append(out, cur.String())
	}
	return out
}

// bitsOf reads a key's size from its blob: the RSA modulus, the curve, or ed25519's fixed 256.
func bitsOf(kind string, blob []byte) int {
	strs := [][]byte{}
	for len(blob) >= 4 {
		n := binary.BigEndian.Uint32(blob)
		if int(n) > len(blob)-4 {
			break
		}
		strs = append(strs, blob[4:4+n])
		blob = blob[4+n:]
	}
	switch {
	case strings.Contains(kind, "ed25519"):
		return 256
	case kind == "ssh-rsa" && len(strs) >= 3:
		return new(big.Int).SetBytes(strs[2]).BitLen()
	case kind == "ssh-dss" && len(strs) >= 2:
		return new(big.Int).SetBytes(strs[1]).BitLen()
	case strings.Contains(kind, "nistp256"):
		return 256
	case strings.Contains(kind, "nistp384"):
		return 384
	case strings.Contains(kind, "nistp521"):
		return 521
	}
	return 0
}