/*
   SSHSecure - a program to harden OpenSSH from defaults
   Copyright (C) 2020  Brent Saner

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

package sshkeys

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"errors"
	"fmt"

	// `golang.org/x/crypto/ssh/internal/bcrypt_pbkdf` Golang doesn't let you import "internal" libs. Fine. Lame language.
	"github.com/dchest/bcrypt_pbkdf"
)

func (k *EncryptedSSHKeyV1) validate() error {
	if k.Passphrase == nil {
		return errors.New("cannot use encrypted key with empty passphrase")
	}
	var validCipher bool
	var validKDF bool
	var validKT bool
	for _, v := range allowed_ciphers {
		if v == k.CipherName {
			validCipher = true
			break
		}
	}
	for _, v := range allowed_kdfnames {
		if v == k.KDFName {
			validKDF = true
			break
		}
	}
	for _, v := range allowed_keytypes {
		if v == k.DefKeyType {
			validKT = true
		}
	}
	if !validCipher || !validKDF || !validKT {
		return errors.New("invalid CipherName, KDFName, or DefKeyType specified")
	}
	return nil
}

func (k *EncryptedSSHKeyV1) Generate(force bool) error {
	if k.DefKeyType == "" {
		k.DefKeyType = defKeyType
	}
	if k.KDFName == "" {
		k.KDFName = defKDF
	}
	if k.CipherName == "" {
		k.CipherName = defCipher
	}
	if err := k.validate(); err != nil {
		return err
	}
	if len(k.Keys) > 0 && !force {
		return nil // Already generated.
	}
	if k.KDFOpts.Salt == nil {
		k.KDFOpts.Salt = make([]byte, defSaltLen)
		if _, err := rand.Read(k.KDFOpts.Salt); err != nil {
			return err
		}
	}
	if k.KDFOpts.Rounds == 0 {
		k.KDFOpts.Rounds = defRounds
	}
	if k.DefKeyType == KeyEd25519 {
		k.KeySize = keyEd25519
		k.BlockSize = blockEd25519
	}
	// Currently, OpenSSH has an option for multiple private keys. However, it is hardcoded to 1.
	// If multiple key support is added in the future, will need to re-tool how I do this, perhaps, in the future. TODO.
	pk := SSHPrivKey{
		Comment: fmt.Sprintf("Autogenerated via SSHSecure (%v)", projUrl),
	}
	pk.Checksum = make([]byte, 4)
	if _, err := rand.Read(pk.Checksum); err != nil {
		return err
	}
	switch k.DefKeyType {
	case KeyRsa:
		if err := pk.generateRsa(); err != nil {
			return err
		}
	case KeyEd25519:
		if err := pk.generateEd25519(); err != nil {
			return err
		}
	default:
		return errors.New("unknown key type; could not generate private/public keypair")
	}
	k.Keys = append(k.Keys, pk)
	// We also need an encrypter/decrypter since this is an encrypted key.
	// Upstream only currently supports bcrypt_pbkdf ("bcrypt").
	// This should always eval to true, but is here for future planning in case other KDF are implemented.
	switch k.KDFName {
	case KdfBcrypt:
		if k.Crypt.CryptKey, err = bcrypt_pbkdf.Key(k.Passphrase, k.KDFOpts.Salt, int(k.KDFOpts.Rounds), kdfKeyLen+len(k.KDFOpts.Salt)); err != nil {
			return err
		} else {
			k.Crypt.PrivateKey = k.Crypt.CryptKey[0:kdfSplit]
			k.Crypt.CryptSalt = k.Crypt.CryptKey[kdfSplit:]
		}
	default:
		return errors.New("could not find KDF")
	}
	switch k.CipherName {
	case CipherAes256Ctr:
		if k.Crypt.Cipher, err = aes.NewCipher(k.Crypt.PrivateKey); err != nil {
			return err
		} else {
			k.Crypt.Stream = cipher.NewCTR(k.Crypt.Cipher, k.Crypt.CryptSalt)
			// Can then be used as k.Crypt.Stream.XORKeyStream(dst []byte, src []byte)
		}
	default:
		return errors.New("could not find Cipher")
	}
	k.build()
	return nil
}

func (k *SSHKeyV1) validate() error {
	var validKT bool
	for _, v := range allowed_keytypes {
		if v == k.DefKeyType {
			validKT = true
		}
	}
	if !validKT {
		return errors.New("invalid DefKeyType specified")
	}
	return nil
}

func (k *SSHKeyV1) Generate(force bool) error {
	if len(k.Keys) > 0 && !force {
		return nil // Already generated.
	}
	if k.DefKeyType == KeyEd25519 {
		k.KeySize = keyEd25519
		k.BlockSize = blockEd25519
	}
	k.CipherName = CipherNull
	k.KDFName = KdfNull
	// Currently, OpenSSH has an option for multiple private keys. However, it is hardcoded to 1.
	// If multiple key support is added in the future, will need to re-tool how I do this, perhaps, in the future. TODO.
	pk := SSHPrivKey{
		Comment: fmt.Sprintf("Autogenerated via SSHSecure (%v)", projUrl),
	}
	pk.Checksum = make([]byte, 4)
	if _, err := rand.Read(pk.Checksum); err != nil {
		return err
	}
	switch k.DefKeyType {
	case KeyRsa:
		if err := pk.generateRsa(); err != nil {
			return err
		}
	case KeyEd25519:
		if err := pk.generateEd25519(); err != nil {
			return err
		}
	default:
		return errors.New("unknown key type; could not generate private/public keypair")
	}
	k.Keys = append(k.Keys, pk)
	k.build()
	return nil
}

func (k *SSHKeyV1) build() error {
	// We actually assemble the key buffer here. Translation to bytes where needed, case switches (ED25519 vs. RSA), etc.
	k.Buffer.Truncate(0)
	// First we need to do some prep for the common header.
	kdfOptsBytes := []byte{}
	switch k.(type) {
	case EncryptedSSHKeyV1:
		kdfOptsBytes = append(kdfOptsBytes, byte(len(k.KDFOpts.Salt)))
	}
	// For this, we can use a generic list and sequentially write.
	cipher := []byte(k.CipherName)
	kdf := []byte(k.KDFName)
	commonHeader := [][]byte{
		[]byte(KeyV1Magic + "\x00"),
		{byte(len(cipher))},
		cipher,
		{byte(len(kdf))},
		kdf,
	}
	for _, v := range commonHeader {
		if _, err := k.Buffer.Write(v); err != nil {
			k.Buffer.Truncate(0)
			return err
		}
	}

}