// License: GPLv3 Copyright: 2022, Kovid Goyal, <kovid at kovidgoyal.net>
// TODO: Move to crypto/ecdh with go >= 1.20

package crypto

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/ecdh"
	"crypto/rand"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"github.com/jamesruan/go-rfc1924/base85"
	"kitty/tools/utils"
	"time"
)

func curve25519_key_pair() (private_key []byte, public_key []byte, err error) {
	curve := ecdh.X25519()
	privkey, err := curve.GenerateKey(rand.Reader)
	if err == nil {
		pubkey := privkey.PublicKey()
		return privkey.Bytes(), pubkey.Bytes(), nil
	}
	return nil, nil, err
}

func curve25519_derive_shared_secret(private_key []byte, public_key []byte) (secret []byte, err error) {
	prkey, err := ecdh.X25519().NewPrivateKey(private_key)
	if err != nil {
		return nil, fmt.Errorf("Invalid X25519 private key: %w", err)
	}
	pubkey, err := ecdh.X25519().NewPublicKey(public_key)
	if err != nil {
		return nil, fmt.Errorf("Invalid X25519 public key: %w", err)
	}
	secret, err = prkey.ECDH(pubkey)
	if err != nil {
		err = fmt.Errorf("Failed to perform ECDH shared secret derivation: %w", err)
	}
	return
}

func b85_encode(data []byte) (encoded string) {
	encoded = base85.EncodeToString(data)
	return
}

func b85_decode(data string) (decoded []byte, err error) {
	decoded, err = base85.DecodeString(data)
	return
}

func encrypt(plaintext []byte, alice_public_key []byte, encryption_protocol string) (iv []byte, tag []byte, ciphertext []byte, bob_public_key []byte, err error) {
	bob_private_key, bob_public_key, err := KeyPair(encryption_protocol)
	if err != nil {
		return
	}
	shared_secret_raw, err := curve25519_derive_shared_secret(bob_private_key, alice_public_key)
	if err != nil {
		return
	}
	shared_secret_hashed := sha256.Sum256(shared_secret_raw)
	shared_secret := shared_secret_hashed[:]
	block, err := aes.NewCipher(shared_secret)
	if err != nil {
		return
	}
	aesgcm, err := cipher.NewGCM(block)
	if err != nil {
		return
	}
	iv = make([]byte, aesgcm.NonceSize())
	_, err = rand.Read(iv)
	if err != nil {
		return
	}
	output := aesgcm.Seal(nil, iv, plaintext, nil)
	ciphertext = output[0 : len(output)-16]
	tag = output[len(output)-16:]
	return
}

func KeyPair(encryption_protocol string) (private_key []byte, public_key []byte, err error) {
	switch encryption_protocol {
	case "1":
		return curve25519_key_pair()
	default:
		err = fmt.Errorf("Unknown encryption protocol: %s", encryption_protocol)
		return
	}
}

func EncodePublicKey(pubkey []byte, encryption_protocol string) (ans string, err error) {
	switch encryption_protocol {
	case "1":
		ans = fmt.Sprintf("1:%s", b85_encode(pubkey))
	default:
		err = fmt.Errorf("Unknown encryption protocol: %s", encryption_protocol)
		return
	}
	return
}

func Encrypt_cmd(cmd *utils.RemoteControlCmd, password string, other_pubkey []byte, encryption_protocol string) (encrypted_cmd utils.EncryptedRemoteControlCmd, err error) {
	cmd.Password = password
	cmd.Timestamp = time.Now().UnixNano()
	plaintext, err := json.Marshal(cmd)
	if err != nil {
		return
	}
	iv, tag, ciphertext, pubkey, err := encrypt(plaintext, other_pubkey, encryption_protocol)
	encrypted_cmd = utils.EncryptedRemoteControlCmd{
		Version: cmd.Version, IV: b85_encode(iv), Tag: b85_encode(tag), Pubkey: b85_encode(pubkey), Encrypted: b85_encode(ciphertext)}
	if encryption_protocol != "1" {
		encrypted_cmd.EncProto = encryption_protocol
	}
	return
}

// }}}