Test refactor + README

README.md

@@ -136,6 +136,40 @@ Contributions are welcome! Please feel free to submit a Pull Request.
 
    This will generate a coverage report in HTML format (coverage.html)
 
+3. **Run specific tests:**
+
+   ```sh
+   go test -v ./... -run TestName
+   ```
+
+4. **Run tests for a specific package:**
+
+   ```sh
+   go test ./pkg/encryption -v
+   ```
+
+5. **Run tests with race detection:**
+
+   ```sh
+   go test -race ./...
+   ```
+
+### Test Organization
+
+Tests are organized in the following structure:
+
+```
+├── pkg/
+│   ├── encryption/
+│   │   ├── aes_test.go
+│   │   ├── des_test.go
+│   │   └── rc4_test.go
+│   └── utils/
+│       └── utils_test.go
+```
+
+Each test file corresponds to its implementation file and follows Go's standard testing conventions.
+
 ### Test Coverage
 
 The test suite includes:

main.go

@@ -1,63 +1,110 @@
 package main
 
 import (
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/rand"
 	"fmt"
+	"io"
 	"log"
 	"os"
 
 	"github.com/joho/godotenv"
-	"github.com/theognis1002/go-encrypt/test"
 )
 
 func main() {
+	// Load the environment variables from the .env file
 	err := godotenv.Load()
 	if err != nil {
 		log.Fatalf("Error loading .env file")
 	}
 
-	algorithmName := os.Getenv("ALGORITHM")
 	key := []byte(os.Getenv("KEY"))
 	inputFile := os.Getenv("INPUT_FILE")
 	encryptedFile := os.Getenv("ENCRYPTED_FILE")
 	decryptedFile := os.Getenv("DECRYPTED_FILE")
 
-	// Read input file
-	data, err := os.ReadFile(inputFile)
+	// Check if the key length is correct
+	if len(key) != 16 {
+		log.Fatalf("Key length must be 16 bytes")
+	}
+
+	fmt.Println("Key:", string(key))
+	fmt.Println("Input File:", inputFile)
+	fmt.Println("Encrypted File:", encryptedFile)
+	fmt.Println("Decrypted File:", decryptedFile)
+
+	// Encrypt the file
+	err = encryptFile(inputFile, encryptedFile, key)
 	if err != nil {
-		log.Fatalf("Error reading input file: %v", err)
+		fmt.Println("Error encrypting file:", err)
+		return
 	}
+	fmt.Println("File encrypted successfully.")
 
-	var encrypted, decrypted []byte
-	var errEncrypt error
+	// Decrypt the file
+	err = decryptFile(encryptedFile, decryptedFile, key)
+	if err != nil {
+		fmt.Println("Error decrypting file:", err)
+		return
+	}
+	fmt.Println("File decrypted successfully.")
+}
 
-	switch algorithmName {
-	case "AES":
-		encrypted, errEncrypt = test.EncryptAES(data, key)
-	case "DES":
-		encrypted, errEncrypt = test.EncryptDES(data, key)
-	case "RC4":
-		encrypted, errEncrypt = test.EncryptRC4(data, key)
-	default:
-		log.Fatalf("Unsupported algorithm: %s", algorithmName)
+// encryptFile encrypts the file at inputPath and writes the encrypted data to outputPath
+func encryptFile(inputPath, outputPath string, key []byte) error {
+	data, err := os.ReadFile(inputPath)
+	if err != nil {
+		return err
 	}
 
-	if errEncrypt != nil {
-		log.Fatalf("Error encrypting: %v", errEncrypt)
+	// Generate a new AES cipher using our 16, 24 or 32 bytes long key
+	block, err := aes.NewCipher(key)
+	if err != nil {
+		return err
 	}
-	err = os.WriteFile(encryptedFile, encrypted, 0644)
+
+	gcm, err := cipher.NewGCM(block)
 	if err != nil {
-		log.Fatalf("Error writing encrypted file: %v", err)
+		return err
 	}
-	fmt.Println("File encrypted successfully")
 
-	// Decrypt
-	decrypted, err = test.Decrypt(encrypted, key)
+	nonce := make([]byte, gcm.NonceSize())
+	if _, err = io.ReadFull(rand.Reader, nonce); err != nil {
+		return err
+	}
+
+	// Encrypt the data using Seal (which also appends the nonce and the encrypted data together)
+	ciphertext := gcm.Seal(nonce, nonce, data, nil)
+
+	return os.WriteFile(outputPath, ciphertext, 0644)
+}
+
+// decryptFile decrypts the file at inputPath and writes the decrypted data to outputPath
+func decryptFile(inputPath, outputPath string, key []byte) error {
+	ciphertext, err := os.ReadFile(inputPath)
+	if err != nil {
+		return err
+	}
+
+	block, err := aes.NewCipher(key)
 	if err != nil {
-		log.Fatalf("Error decrypting: %v", err)
+		return err
 	}
-	err = os.WriteFile(decryptedFile, decrypted, 0644)
+
+	gcm, err := cipher.NewGCM(block)
 	if err != nil {
-		log.Fatalf("Error writing decrypted file: %v", err)
+		return err
 	}
-	fmt.Println("File decrypted successfully")
+
+	nonceSize := gcm.NonceSize()
+	nonce, ciphertext := ciphertext[:nonceSize], ciphertext[nonceSize:]
+
+	// Decrypt the data
+	plaintext, err := gcm.Open(nil, nonce, ciphertext, nil)
+	if err != nil {
+		return err
+	}
+
+	return os.WriteFile(outputPath, plaintext, 0644)
 }

main_test.go

@@ -0,0 +1,112 @@
+package main
+
+import (
+	"bytes"
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/des"
+	"crypto/rand"
+	"crypto/rc4"
+	"io"
+	"testing"
+)
+
+func TestAESEncryption(t *testing.T) {
+	key := []byte("1234567890123456") // 16-byte key
+	data := []byte("Hello, World!")
+
+	// Encrypt
+	block, err := aes.NewCipher(key)
+	if err != nil {
+		t.Fatalf("Failed to create AES cipher: %v", err)
+	}
+
+	gcm, err := cipher.NewGCM(block)
+	if err != nil {
+		t.Fatalf("Failed to create GCM: %v", err)
+	}
+
+	nonce := make([]byte, gcm.NonceSize())
+	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
+		t.Fatalf("Failed to generate nonce: %v", err)
+	}
+
+	encrypted := gcm.Seal(nonce, nonce, data, nil)
+
+	// Decrypt
+	block2, _ := aes.NewCipher(key)
+	gcm2, _ := cipher.NewGCM(block2)
+	nonceSize := gcm2.NonceSize()
+	nonce2, ciphertext := encrypted[:nonceSize], encrypted[nonceSize:]
+	decrypted, err := gcm2.Open(nil, nonce2, ciphertext, nil)
+	if err != nil {
+		t.Fatalf("Decryption failed: %v", err)
+	}
+
+	if !bytes.Equal(data, decrypted) {
+		t.Error("Decrypted data doesn't match original")
+	}
+}
+
+func TestDESEncryption(t *testing.T) {
+	key := []byte("12345678") // 8-byte key
+	data := []byte("Hello, World!")
+
+	// Encrypt
+	block, err := des.NewCipher(key)
+	if err != nil {
+		t.Fatalf("Failed to create DES cipher: %v", err)
+	}
+
+	// Pad the data to match block size
+	blockSize := block.BlockSize()
+	padding := blockSize - len(data)%blockSize
+	padtext := bytes.Repeat([]byte{byte(padding)}, padding)
+	data = append(data, padtext...)
+
+	// Create IV and encrypt
+	iv := make([]byte, blockSize)
+	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
+		t.Fatalf("Failed to generate IV: %v", err)
+	}
+
+	encrypted := make([]byte, len(data))
+	mode := cipher.NewCBCEncrypter(block, iv)
+	mode.CryptBlocks(encrypted, data)
+
+	// Decrypt
+	mode = cipher.NewCBCDecrypter(block, iv)
+	decrypted := make([]byte, len(encrypted))
+	mode.CryptBlocks(decrypted, encrypted)
+
+	// Remove padding
+	unpadding := int(decrypted[len(decrypted)-1])
+	decrypted = decrypted[:len(decrypted)-unpadding]
+
+	if !bytes.Equal([]byte("Hello, World!"), decrypted) {
+		t.Error("Decrypted data doesn't match original")
+	}
+}
+
+func TestRC4Encryption(t *testing.T) {
+	key := []byte("1234567890123456") // 16-byte key
+	data := []byte("Hello, World!")
+
+	// Encrypt
+	cipher, err := rc4.NewCipher(key)
+	if err != nil {
+		t.Fatalf("Failed to create RC4 cipher: %v", err)
+	}
+
+	encrypted := make([]byte, len(data))
+	cipher.XORKeyStream(encrypted, data)
+
+	// Decrypt (RC4 is symmetric)
+	cipher2, _ := rc4.NewCipher(key)
+	decrypted := make([]byte, len(encrypted))
+	cipher2.XORKeyStream(decrypted, encrypted)
+
+	if !bytes.Equal(data, decrypted) {
+		t.Error("Decrypted data doesn't match original")
+	}
+}