Encrypt and decrypt an archive

archive/archive.go

@@ -1,6 +1,9 @@
 package archive
 
-import "io"
+import (
+	"bytes"
+	"io"
+)
 
 // An Archive represents a collection of files stored in a single file.
 type Archive struct {
@@ -20,6 +23,16 @@ func (a *Archive) TotalSize() uint64 {
 	return total
 }
 
+// GetBytes returns the archive as a byte slice.
+func (a *Archive) GetBytes() ([]byte, error) {
+	data := new(bytes.Buffer)
+	if err := a.Write(data); err != nil {
+		return nil, err
+	}
+
+	return data.Bytes(), nil
+}
+
 // Write writes the archive into the provided writer.
 func (a *Archive) Write(w io.Writer) error {
 	if err := a.Header.Write(w); err != nil {

archive/encryption.go

@@ -0,0 +1,144 @@
+package archive
+
+import (
+	"bytes"
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/rand"
+	"crypto/sha256"
+	"io"
+
+	"golang.org/x/crypto/pbkdf2"
+)
+
+// An EncryptedArchive represents an encrypted archive.
+type EncryptedArchive struct {
+	bytes []byte
+	salt  []byte
+	nonce []byte
+}
+
+// Write writes the encrypted archive into the provided writer.
+// The encrypted archive is serialized as follows:
+//
+//  1. The magic field is serialized as a 4-byte sequence.
+//  2. The salt field is serialized as a 16-byte sequence.
+//  3. The nonce field is serialized as a sequence of bytes.
+//  4. The encrypted data is serialized as a sequence of bytes.
+func (a *EncryptedArchive) Write(w io.Writer) error {
+	// Write the magic (4 bytes)
+	if _, err := w.Write(encMagic); err != nil {
+		return err
+	}
+
+	// Write the salt (16 bytes)
+	if _, err := w.Write(a.salt); err != nil {
+		return err
+	}
+
+	// Write the nonce
+	if _, err := w.Write(a.nonce); err != nil {
+		return err
+	}
+
+	// Write the encrypted data
+	if _, err := w.Write(a.bytes); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// ReadEncryptedArchive reads an encrypted archive from the provided reader.
+func ReadEncryptedArchive(r io.Reader) (*EncryptedArchive, error) {
+	if err := mustReadEncryptedMagic(r); err != nil {
+		return nil, err
+	}
+
+	// Read the salt (16 bytes)
+	salt := make([]byte, 16)
+	if _, err := io.ReadFull(r, salt); err != nil {
+		return nil, err
+	}
+
+	// Read the nonce
+	nonce := make([]byte, 12)
+	if _, err := io.ReadFull(r, nonce); err != nil {
+		return nil, err
+	}
+
+	// Read the encrypted data
+	data, err := io.ReadAll(r)
+	if err != nil {
+		return nil, err
+	}
+
+	return &EncryptedArchive{
+		bytes: data,
+		salt:  salt,
+		nonce: nonce,
+	}, nil
+}
+
+// Encrypt encrypts the archive using AES-GCM with the provided password.
+func (a *Archive) Encrypt(password string) (*EncryptedArchive, error) {
+	// Generate a salt for key derivation (PBKDF2)
+	salt := make([]byte, 16)
+	if _, err := rand.Read(salt); err != nil {
+		return nil, err
+	}
+
+	aesGCM, err := newCipher(password, salt)
+	if err != nil {
+		return nil, err
+	}
+
+	// Generate a nonce for AES-GCM (random IV)
+	nonce := make([]byte, aesGCM.NonceSize())
+	if _, err := rand.Read(nonce); err != nil {
+		return nil, err
+	}
+
+	// Get the plaintext data
+	plaintext, err := a.GetBytes()
+	if err != nil {
+		return nil, err
+	}
+
+	// Encrypt the data using AES-GCM
+	ciphertext := aesGCM.Seal(nil, nonce, plaintext, nil)
+
+	return &EncryptedArchive{
+		bytes: ciphertext,
+		salt:  salt,
+		nonce: nonce,
+	}, nil
+}
+
+// Decrypt decrypts the encrypted archive using AES-GCM with the provided password.
+// If the password is incorrect, the process will fail as the Archive data will be
+// corrupted.
+func (a *EncryptedArchive) Decrypt(password string) (*Archive, error) {
+	aesGCM, err := newCipher(password, a.salt)
+	if err != nil {
+		return nil, err
+	}
+
+	plaintext, err := aesGCM.Open(nil, a.nonce, a.bytes, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	return ReadArchive(bytes.NewReader(plaintext))
+}
+
+// newCipher creates a new AES-GCM cipher with the provided password and salt.
+func newCipher(password string, salt []byte) (cipher.AEAD, error) {
+	key := pbkdf2.Key([]byte(password), salt, 4096, 32, sha256.New)
+	block, err := aes.NewCipher(key)
+	if err != nil {
+		return nil, err
+	}
+
+	return cipher.NewGCM(block)
+}

archive/encryption_test.go

@@ -0,0 +1,76 @@
+package archive
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestEncryptDecryptArchive(t *testing.T) {
+	archive := makeTestArchive()
+
+	encrypted, err := archive.Encrypt("password")
+	assert.Nil(t, err)
+
+	decrypted, err := encrypted.Decrypt("password")
+	assert.Nil(t, err)
+
+	assert.Equal(t, archive, decrypted)
+}
+
+func TestWriteAndReadEncryptedArchive(t *testing.T) {
+	var (
+		archive      = makeTestArchive()
+		encrypted, _ = archive.Encrypt("password")
+		w            = new(bytes.Buffer)
+	)
+
+	err := encrypted.Write(w)
+	assert.Nil(t, err)
+
+	r := bytes.NewReader(w.Bytes())
+	readArchive, err := ReadEncryptedArchive(r)
+	assert.Nil(t, err)
+
+	assert.Equal(t, encrypted, readArchive)
+}
+
+func makeTestArchive() *Archive {
+	return &Archive{
+		Header: &Header{
+			HeaderLength: 46,
+			Entries: []*HeaderFileEntry{
+				{
+					Name:   "file1.txt",
+					Size:   12,
+					Offset: 28,
+				},
+				{
+					Name:   "file2.txt",
+					Size:   16,
+					Offset: 40,
+				},
+			},
+		},
+		Files: []*ArchiveFile{
+			{
+				FileName: "file1.txt",
+				CompressedBytes: []byte{
+					0x78, 0x9c, 0x4b, 0x4c,
+					0x4f, 0x49, 0x2d, 0x2e,
+					0x01, 0x00, 0x00, 0xff,
+				},
+			},
+			{
+				FileName: "file2.txt",
+				CompressedBytes: []byte{
+					0x78, 0x9c, 0x4b, 0x4c,
+					0x4f, 0x49, 0x2d, 0x2e,
+					0x01, 0x00, 0x00, 0xff,
+					0x78, 0x9c, 0x4b, 0x4c,
+				},
+			},
+		},
+	}
+}

archive/header.go

@@ -6,13 +6,6 @@ import (
 	"io"
 )
 
-// magic is a unique identifier for the archive format.
-// It's the ASCII representation of "AAR?".
-var magic = []byte{0x41, 0x41, 0x52, 0x3F}
-
-// magicLen is the length of the magic field in bytes.
-const magicLen = uint32(4)
-
 // byteOrder is the byte order used to serialize integers.
 var byteOrder = binary.LittleEndian
 

archive/magic.go

@@ -0,0 +1,59 @@
+package archive
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+)
+
+// magic is a unique identifier for the archive format.
+// It's the ASCII representation of "AAR?".
+var magic = []byte{0x41, 0x41, 0x52, 0x3F}
+
+// encMagic is a unique identifier for the encrypted archive format.
+// It's the ASCII representation of "AARX".
+var encMagic = []byte{0x41, 0x41, 0x52, 0x58}
+
+// magicLen is the length of the magic field in bytes.
+const magicLen = uint32(4)
+
+// ErrInvalidMagic is returned when the magic field is not correct.
+var ErrInvalidMagic = fmt.Errorf("invalid magic, expected %v", magic)
+
+// ErrInvalidEncMagic is returned when the magic field is not correct.
+var ErrInvalidEncMagic = fmt.Errorf("invalid magic, expected %v", encMagic)
+
+// mustReadMagic reads the magic field from the provided reader.
+// If the magic field is not correct, it returns an error.
+func mustReadMagic(r io.Reader) error {
+	readMagic := make([]byte, 4)
+
+	// Read the magic (4 bytes)
+	if _, err := io.ReadFull(r, readMagic); err != nil {
+		return err
+	}
+
+	// Check if the magic is correct
+	if !bytes.Equal(magic, readMagic) {
+		return ErrInvalidMagic
+	}
+
+	return nil
+}
+
+// mustReadEncryptedMagic reads the magic field from the provided reader.
+func mustReadEncryptedMagic(r io.Reader) error {
+	readMagic := make([]byte, 4)
+
+	// Read the magic (4 bytes)
+	if _, err := io.ReadFull(r, readMagic); err != nil {
+		return err
+	}
+
+	// Check if the magic is correct
+	if !bytes.Equal(encMagic, readMagic) {
+		return ErrInvalidEncMagic
+	}
+
+	return nil
+}

go.mod

@@ -5,6 +5,7 @@ go 1.22.0
 require (
 	github.com/stretchr/testify v1.9.0
 	github.com/ulikunitz/xz v0.5.12
+	golang.org/x/crypto v0.26.0
 )
 
 require (

go.sum

@@ -6,6 +6,8 @@ github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsT
 github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
 github.com/ulikunitz/xz v0.5.12 h1:37Nm15o69RwBkXM0J6A5OlE67RZTfzUxTj8fB3dfcsc=
 github.com/ulikunitz/xz v0.5.12/go.mod h1:nbz6k7qbPmH4IRqmfOplQw/tblSgqTqBwxkY0oWt/14=
+golang.org/x/crypto v0.26.0 h1:RrRspgV4mU+YwB4FYnuBoKsUapNIL5cohGAmSH3azsw=
+golang.org/x/crypto v0.26.0/go.mod h1:GY7jblb9wI+FOo5y8/S2oY4zWP07AkOJ4+jxCqdqn54=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=