package sshclient

import (
	"crypto/aes"
	"crypto/cipher"
	"encoding/pem"
	"errors"
	"fmt"

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

// aes256-gcm@openssh.com is a cipher OpenSSH (and 1Password key exports) may
// use for private key files. golang.org/x/crypto/ssh only supports aes256-ctr
// and aes256-cbc, so these keys are decrypted here and re-marshalled as
// unencrypted openssh-key-v1 blobs that ssh.ParsePrivateKey accepts.
const (
	opensshKeyMagic  = "openssh-key-v1\x00"
	opensshGCMCipher = "aes256-gcm@openssh.com"
	gcmTagSize       = 16
)

type opensshPrivateKeyHeader struct {
	CipherName   string
	KdfName      string
	KdfOpts      string
	NumKeys      uint32
	PubKey       []byte
	PrivKeyBlock []byte
	// For AEAD ciphers (aes256-gcm@openssh.com) OpenSSH stores the 16-byte
	// auth tag as raw bytes after the encrypted block, not inside it.
	Rest []byte `ssh:"rest"`
}

// isGCMEncryptedKey reports whether data is an OpenSSH private key encrypted
// with aes256-gcm@openssh.com.
func isGCMEncryptedKey(data []byte) bool {
	block, _ := pem.Decode(data)
	if block == nil || block.Type != "OPENSSH PRIVATE KEY" {
		return false
	}
	if len(block.Bytes) < len(opensshKeyMagic) || string(block.Bytes[:len(opensshKeyMagic)]) != opensshKeyMagic {
		return false
	}
	var header struct {
		CipherName string
		Rest       []byte `ssh:"rest"`
	}
	if err := ssh.Unmarshal(block.Bytes[len(opensshKeyMagic):], &header); err != nil {
		return false
	}
	return header.CipherName == opensshGCMCipher
}

// parseGCMKey decrypts an aes256-gcm@openssh.com encrypted OpenSSH private
// key with passphrase and returns the signer.
func parseGCMKey(data, passphrase []byte) (ssh.Signer, error) {
	block, _ := pem.Decode(data)
	if block == nil || len(block.Bytes) < len(opensshKeyMagic) {
		return nil, errors.New("invalid OpenSSH private key PEM")
	}
	var w opensshPrivateKeyHeader
	if err := ssh.Unmarshal(block.Bytes[len(opensshKeyMagic):], &w); err != nil {
		return nil, err
	}
	if w.NumKeys != 1 {
		return nil, errors.New("multi-key files are not supported")
	}
	if w.KdfName != "bcrypt" {
		return nil, fmt.Errorf("unknown KDF %q, only supports %q", w.KdfName, "bcrypt")
	}
	var kdfOpts struct {
		Salt   string
		Rounds uint32
	}
	if err := ssh.Unmarshal([]byte(w.KdfOpts), &kdfOpts); err != nil {
		return nil, err
	}
	// Same cap as x/crypto/ssh so a hostile file cannot tie up the caller.
	const maxRounds = 1 << 11
	if kdfOpts.Rounds > maxRounds {
		return nil, fmt.Errorf("bcrypt KDF rounds %d exceed maximum %d", kdfOpts.Rounds, maxRounds)
	}

	// AES-256-GCM: 32-byte key + 12-byte nonce derived by bcrypt_pbkdf.
	// The 16-byte auth tag follows the encrypted block as raw bytes;
	// gcm.Open expects it appended to the ciphertext.
	if len(w.Rest) != gcmTagSize {
		return nil, fmt.Errorf("unexpected GCM tag length %d", len(w.Rest))
	}
	k, err := bcryptPBKDF(passphrase, []byte(kdfOpts.Salt), int(kdfOpts.Rounds), 32+12)
	if err != nil {
		return nil, err
	}
	aesCipher, err := aes.NewCipher(k[:32])
	if err != nil {
		return nil, err
	}
	gcm, err := cipher.NewGCM(aesCipher)
	if err != nil {
		return nil, err
	}
	sealed := make([]byte, 0, len(w.PrivKeyBlock)+gcmTagSize)
	sealed = append(sealed, w.PrivKeyBlock...)
	sealed = append(sealed, w.Rest...)
	plain, err := gcm.Open(nil, k[32:], sealed, nil)
	if err != nil {
		return nil, err
	}

	// Re-marshal as an unencrypted key (tag dropped). The plaintext keeps
	// its original OpenSSH padding, which ssh.ParsePrivateKey validates.
	w.CipherName = "none"
	w.KdfName = "none"
	w.KdfOpts = ""
	w.PrivKeyBlock = plain
	w.Rest = nil
	raw := make([]byte, 0, len(opensshKeyMagic)+len(plain)+128)
	raw = append(raw, opensshKeyMagic...)
	raw = append(raw, ssh.Marshal(&w)...)
	return ssh.ParsePrivateKey(pem.EncodeToMemory(&pem.Block{Type: "OPENSSH PRIVATE KEY", Bytes: raw}))
}