package main

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/pem"
	"errors"
	"fmt"
	"os"
    "flag"
    
    "golang.org/x/crypto/ssh/terminal"
	"golang.org/x/crypto/scrypt"
)

func main() {
  	keyPtr := flag.String("keyfile", "encrypted_aes_key.pem", "Enter keyfile")
	flag.Parse()
    
    if *keyPtr == "" {
        fmt.Print("Please enter a keyfile")
        return
    }
    
	// Generate a new AES-256 key
	key, err := generateAESKey(32) // 32 bytes = 256 bits
	if err != nil {
		fmt.Printf("Error generating AES key: %v\n", err)
		return
	}

	// Prompt for passphrase
	print("Enter passphrase: ")
	pass, _ := terminal.ReadPassword(0) // Hides input
    password := string(pass)

	// Encrypt and PEM encode the key
	pemBlock, err := encryptKeyToPEM(key, password)
	if err != nil {
		fmt.Printf("Error encrypting key: %v\n", err)
		return
	}

	// Save to file
	err = os.WriteFile(*keyPtr, pem.EncodeToMemory(pemBlock), 0600)
	if err != nil {
		fmt.Printf("Error writing PEM file: %v\n", err)
		return
	}

	fmt.Printf("Successfully generated and password-protected AES key in %s", *keyPtr)
}

// generateAESKey creates a new random AES key of specified size
func generateAESKey(size int) ([]byte, error) {
	key := make([]byte, size)
	_, err := rand.Read(key)
	if err != nil {
		return nil, err
	}
	return key, nil
}

// encryptKeyToPEM encrypts the key with a password and returns a PEM block
func encryptKeyToPEM(key []byte, password string) (*pem.Block, error) {
	// Convert password to suitable encryption key
	salt := make([]byte, 16)
	if _, err := rand.Read(salt); err != nil {
		return nil, err
	}

	// Use scrypt for key derivation (better than PBKDF2)
	// N: CPU/memory cost, r: block size, p: parallelization
	encryptionKey, err := scrypt.Key([]byte(password), salt, 32768, 8, 1, 32)
	if err != nil {
		return nil, err
	}

	// Generate a random IV
	iv := make([]byte, aes.BlockSize)
	if _, err := rand.Read(iv); err != nil {
		return nil, err
	}

	// Encrypt the key
	block, err := aes.NewCipher(encryptionKey)
	if err != nil {
		return nil, err
	}

	encrypted := make([]byte, len(key))
	stream := cipher.NewCFBEncrypter(block, iv)
	stream.XORKeyStream(encrypted, key)

	// Create PEM block
	pemBlock := &pem.Block{
		Type: "ENCRYPTED AES KEY",
		Headers: map[string]string{
			"Key-Derivation": "scrypt",
			"SALT":           fmt.Sprintf("%x", salt),
			"IV":             fmt.Sprintf("%x", iv),
			"Key-Size":       fmt.Sprintf("%d", len(key)*8),
		},
		Bytes: encrypted,
	}

	return pemBlock, nil
}

// hexStringToBytes converts a hex string to bytes
func hexStringToBytes(hexStr string) ([]byte, error) {
	if len(hexStr)%2 != 0 {
		return nil, errors.New("hex string length must be even")
	}
	bytes := make([]byte, len(hexStr)/2)
	for i := 0; i < len(hexStr); i += 2 {
		b, err := hexByte(hexStr[i], hexStr[i+1])
		if err != nil {
			return nil, err
		}
		bytes[i/2] = b
	}
	return bytes, nil
}

// hexByte converts two hex characters to a byte
func hexByte(a, b byte) (byte, error) {
	high, err := hexDigitToByte(a)
	if err != nil {
		return 0, err
	}
	low, err := hexDigitToByte(b)
	if err != nil {
		return 0, err
	}
	return (high << 4) | low, nil
}

// hexDigitToByte converts a single hex character to its byte value
func hexDigitToByte(c byte) (byte, error) {
	switch {
	case '0' <= c && c <= '9':
		return c - '0', nil
	case 'a' <= c && c <= 'f':
		return c - 'a' + 10, nil
	case 'A' <= c && c <= 'F':
		return c - 'A' + 10, nil
	default:
		return 0, fmt.Errorf("invalid hex digit %c", c)
	}
}