package coldstorage

import (
	"crypto/rand"
	"crypto/rsa"
	"errors"
	"fmt"
	"os"
	"path"
	"path/filepath"

	"code.nonshy.com/nonshy/website/pkg/encryption/keygen"
	"code.nonshy.com/nonshy/website/pkg/log"
)

var (
	ColdStorageDirectory      = "./coldstorage"
	ColdStorageKeysDirectory  = path.Join(ColdStorageDirectory, "keys")
	ColdStoragePrivateKeyFile = path.Join(ColdStorageKeysDirectory, "private.pem")
	ColdStoragePublicKeyFile  = path.Join(ColdStorageKeysDirectory, "public.pem")
)

// Initialize generates the RSA key pairs for the first time and creates
// the cold storage directories. It writes the keys to disk and returns the x509 encoded
// public key which goes in the settings.json (the keys on disk are for your bookkeeping).
func Initialize() ([]byte, error) {
	log.Warn("NOTICE: rolling a random RSA key pair for cold storage")
	rsaKey, err := keygen.NewRSAKeys()
	if err != nil {
		return nil, fmt.Errorf("generate RSA key: %s", err)
	}

	// Encode to x509
	x509, err := keygen.SerializePublicKey(rsaKey.Public())
	if err != nil {
		return nil, fmt.Errorf("encode RSA public key to x509: %s", err)
	}

	// Write the public/private key files to disk.
	if _, err := os.Stat(ColdStorageKeysDirectory); os.IsNotExist(err) {
		log.Info("Notice: creating cold storage directory")
		if err := os.MkdirAll(ColdStorageKeysDirectory, 0755); err != nil {
			return nil, fmt.Errorf("create %s: %s", ColdStorageKeysDirectory, err)
		}
	}
	if err := keygen.WriteRSAKeys(
		rsaKey,
		path.Join(ColdStorageKeysDirectory, "private.pem"),
		path.Join(ColdStorageKeysDirectory, "public.pem"),
	); err != nil {
		return nil, fmt.Errorf("export newly generated public/private key files: %s", err)
	}

	return x509, nil
}

// Warning returns an error message if the private key is still on disk at its
// original generated location: it should be moved offline for security.
func Warning() error {
	if _, err := os.Stat(ColdStoragePrivateKeyFile); os.IsNotExist(err) {
		return nil
	}

	return errors.New("the private key file at ./coldstorage/keys should be moved off of the server and kept offline for safety")
}

// FileToColdStorage will copy a file, encrypted, into cold storage at the given filename.
func FileToColdStorage(sourceFilePath, outputFileName string, publicKeyPEM []byte) error {
	if len(publicKeyPEM) == 0 {
		return errors.New("no RSA public key")
	}

	// Load the public key from PEM encoding.
	publicKey, err := keygen.DeserializePublicKey(publicKeyPEM)
	if err != nil {
		return fmt.Errorf("deserializing public key: %s", err)
	}

	// Generate a unique AES key for encrypting this file in one direction.
	aesKey, err := keygen.NewAESKey()
	if err != nil {
		return err
	}

	// Encrypt the AES key and store it on disk next to the cold storage file.
	ciphertext, err := rsa.EncryptPKCS1v15(rand.Reader, publicKey, aesKey)
	if err != nil {
		return fmt.Errorf("encrypt error: %s", err)
	}
	err = os.WriteFile(
		filepath.Join(ColdStorageDirectory, outputFileName+".aes"),
		ciphertext,
		0600,
	)
	if err != nil {
		return err
	}

	// Read the original plaintext file going into cold storage.
	plaintext, err := os.ReadFile(sourceFilePath)
	if err != nil {
		return fmt.Errorf("source file: %s", err)
	}

	// Encrypt the plaintext with the AES key.
	ciphertext, err = keygen.EncryptWithAESKey(plaintext, aesKey)
	if err != nil {
		return err
	}

	// Write it to disk.
	return os.WriteFile(filepath.Join(ColdStorageDirectory, outputFileName+".enc"), ciphertext, 0600)
}

// FileFromColdStorage decrypts a cold storage file and writes it to the output file.
//
// The command `nonshy coldstorage decrypt` uses this function. Requirements:
//
// - privateKeyFile is the RSA private key originally generated for cold storage
// - aesKeyFile is the unique .aes file for the cold storage item
// - ciphertextFile is the encrypted cold storage item
// - outputFile is where you want to save the result to
func FileFromColdStorage(privateKeyFile, aesKeyFile, ciphertextFile, outputFile string) error {
	privateKey, err := keygen.PrivateKeyFromFile(privateKeyFile)
	if err != nil {
		return fmt.Errorf("private key file: %s", err)
	}

	encryptedAESKey, err := os.ReadFile(aesKeyFile)
	if err != nil {
		return fmt.Errorf("reading aes key file: %s", err)
	}

	aesKey, err := rsa.DecryptPKCS1v15(rand.Reader, privateKey, encryptedAESKey)
	if err != nil {
		return fmt.Errorf("decrypting the aes key file: %s", err)
	}

	ciphertext, err := os.ReadFile(ciphertextFile)
	if err != nil {
		return fmt.Errorf("reading cold storage file: %s", err)
	}

	plaintext, err := keygen.DecryptWithAESKey(ciphertext, aesKey)
	if err != nil {
		return fmt.Errorf("decrypting cold storage file: %s", err)
	}

	return os.WriteFile(outputFile, plaintext, 0644)
}