/*
   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"
	"errors"

	// `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) setCrypt() error {
	// First we need the KDF (Key Derivation Function).
	if err := k.setKDF(); err != nil {
		return err
	}
	// And then we set the actual cipher (the thing that does the encrypting).
	if err := k.setCipher(); err != nil {
		return err
	}
	return nil
}

func (k *EncryptedSSHKeyV1) setCipher() error {
	var err error
	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")
	}
	return nil
}

func (k *EncryptedSSHKeyV1) setKDF() error {
	var err error
	// 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")
	}
	return nil
}