feat: add aes encryption method. (#10)

This commit refactor project to include multiple encryption methods.
Add Aes encryption, loader and tests

Part of: #2

Signed-off-by: Vasek - Tom C <[email protected]>

Author: TomChv

.gitignore

@@ -22,5 +22,6 @@ go.work
 
 221b
 *.sh
+payload_usage/
 
 .idea

cli/bake.go

@@ -7,7 +7,6 @@ import (
 	"path/filepath"
 	"strings"
 
-	"github.com/cmepw/221b/encryption/xor"
 	"github.com/cmepw/221b/loader"
 	"github.com/cmepw/221b/logger"
 )
@@ -16,11 +15,13 @@ var (
 	shellPath string
 	key       string
 	output    string
+	method    string
 )
 
 var (
-	ErrMissingShellPath = fmt.Errorf("missing shellPath argument")
-	ErrMissingKey       = fmt.Errorf("missing key argument")
+	ErrMissingShellPath   = fmt.Errorf("missing shellPath argument")
+	ErrMissingKey         = fmt.Errorf("missing key argument")
+	ErrMethodNotSupported = fmt.Errorf("provided encryption method isn't supported, please choose: 'aes', 'xor'")
 )
 
 var bake = &cobra.Command{
@@ -40,6 +41,12 @@ var bake = &cobra.Command{
 			logger.Warn(fmt.Sprintf("output path set to default: %s", output))
 		}
 
+		loader, ok := loader.Method[method]
+		if !ok {
+			logger.Fatal(ErrMethodNotSupported)
+		}
+
+		logger.Info(fmt.Sprintf("use %s encryption method", method))
 		logger.Debug(fmt.Sprintf("baking %s with key %s", shellPath, key))
 
 		logger.Debug(fmt.Sprintf("reading %s", shellPath))
@@ -48,13 +55,16 @@ var bake = &cobra.Command{
 			logger.Fatal(err)
 		}
 
-		logger.Debug(fmt.Sprintf("encrypting %s", shellPath))
-		encryptedShell := xor.Encrypt(file, []byte(key))
+		logger.Info(fmt.Sprintf("encrypting %s", shellPath))
 
-		logger.Debug(fmt.Sprintf("injecting encrypted shell into payload"))
+		encryptedShell, err := loader.Encrypt(file, []byte(key))
+		if err != nil {
+			logger.Fatal(err)
+		}
+
+		logger.Info(fmt.Sprintf("loading encrypted shell into payload"))
 
-		xorLoader := loader.NewXorLoader([]byte(key))
-		content, err := xorLoader.Load(encryptedShell)
+		content, err := loader.Load(encryptedShell, []byte(key))
 		if err != nil {
 			logger.Fatal(err)
 		}
@@ -64,7 +74,9 @@ var bake = &cobra.Command{
 			fmt.Println(string(content))
 		}
 
-		if err := xorLoader.Compile(output, content); err != nil {
+		logger.Info("compiling binary")
+
+		if err := loader.Compile(output, content); err != nil {
 			logger.Fatal(err)
 		}
 	},
@@ -74,4 +86,5 @@ func init() {
 	bake.Flags().StringVarP(&shellPath, "shellPath", "s", "", "path to the shell scrypt")
 	bake.Flags().StringVarP(&key, "key", "k", "", "key to use for the xor")
 	bake.Flags().StringVarP(&output, "output", "o", "", "output path (e.g., /home/bin.exe)")
+	bake.Flags().StringVarP(&method, "method", "m", "xor", "encryption method")
 }

encryption/aes.go

@@ -0,0 +1,53 @@
+package encryption
+
+import (
+	"crypto/aes"
+	"crypto/cipher"
+	"crypto/rand"
+	"errors"
+	"io"
+)
+
+var ErrCipherTextTooShort = errors.New("ciphertext too short")
+
+type Aes struct{}
+
+func (a Aes) Decrypt(content, key []byte) ([]byte, error) {
+	c, err := aes.NewCipher(key)
+	if err != nil {
+		return nil, err
+	}
+
+	gcm, err := cipher.NewGCM(c)
+	if err != nil {
+		return nil, err
+	}
+
+	nonceSize := gcm.NonceSize()
+	if len(content) < nonceSize {
+		return nil, ErrCipherTextTooShort
+	}
+
+	nonce, ciphertext := content[:nonceSize], content[nonceSize:]
+
+	return gcm.Open(nil, nonce, ciphertext, nil)
+}
+
+func (a Aes) Encrypt(content, key []byte) ([]byte, error) {
+	c, err := aes.NewCipher(key)
+	if err != nil {
+		return nil, err
+	}
+
+	gcm, err := cipher.NewGCM(c)
+	if err != nil {
+		return nil, err
+	}
+
+	nonce := make([]byte, gcm.NonceSize())
+	if _, err = io.ReadFull(rand.Reader, nonce); err != nil {
+		return nil, err
+	}
+
+	return gcm.Seal(nonce, nonce, content, nil), nil
+}

encryption/aes_test.go

@@ -0,0 +1,47 @@
+package encryption
+
+import (
+	"bytes"
+	"github.com/stretchr/testify/require"
+	"testing"
+)
+
+func TestAes_EncryptDecrypt(t *testing.T) {
+	type TestCase struct {
+		content []byte
+		key     []byte
+	}
+
+	testsCases := map[string]TestCase{
+		"encrypt a simple string": {
+			content: []byte("hello world"),
+			key:     []byte("0123456789ABCDEF"),
+		},
+		"encrypt a golang file": {
+			content: []byte(`
+package main
+
+import "fmt"
+
+func main() {
+	fmt.Println("hello Go!")
+}
+`),
+			key: []byte("0123456789ABCDEF"),
+		},
+	}
+
+	for name, tt := range testsCases {
+		t.Run(name, func(t *testing.T) {
+			aes := Aes{}
+
+			encryptedContent, err := aes.Encrypt(bytes.Clone(tt.content), tt.key)
+			require.NoError(t, err)
+			require.NotEqual(t, tt.content, encryptedContent)
+
+			decryptedContent, err := aes.Decrypt(encryptedContent, tt.key)
+			require.NoError(t, err)
+			require.Equal(t, tt.content, decryptedContent)
+		})
+	}
+}

encryption/encryption.go

@@ -1,6 +1,6 @@
 package encryption
 
 type Encryption interface {
-	Decrypt(content, key []byte) []byte
-	Encrypt(content, key []byte) []byte
+	Decrypt(content, key []byte) ([]byte, error)
+	Encrypt(content, key []byte) ([]byte, error)
 }

encryption/xor/xor.go encryption/xor.go

@@ -1,25 +1,27 @@
-// Package xor implements xor encryption and decryption.
-package xor
+package encryption
+
+// Xor implements xor encryption and decryption.
+type Xor struct{}
 
 // Decrypt apply xor on the content with the given key.
 // Note that it will actually update the given content.
-func Decrypt(content, key []byte) []byte {
+func (x Xor) Decrypt(content, key []byte) ([]byte, error) {
 	keyLen := len(key)
 
 	for i := 0; i < len(content); i++ {
 		content[i] ^= key[i%keyLen]
 	}
-	return content
+	return content, nil
 }
 
 // Encrypt apply xor on the content with the given key.
 // Note that it will actually update the given content.
-func Encrypt(content, key []byte) []byte {
+func (x Xor) Encrypt(content, key []byte) ([]byte, error) {
 	keyLen := len(key)
 
 	for i := 0; i < len(content); i++ {
 		content[i] ^= key[i%keyLen]
 	}
 
-	return content
+	return content, nil
 }

encryption/xor/xor_test.go encryption/xor_test.go

@@ -1,4 +1,4 @@
-package xor
+package encryption
 
 import (
 	"bytes"
@@ -33,10 +33,14 @@ func main() {
 
 	for name, tt := range testsCases {
 		t.Run(name, func(t *testing.T) {
-			encryptedContent := Encrypt(bytes.Clone(tt.content), tt.key)
+			xor := Xor{}
 
+			encryptedContent, err := xor.Encrypt(bytes.Clone(tt.content), tt.key)
+			require.NoError(t, err)
 			require.NotEqual(t, tt.content, encryptedContent)
-			decryptedContent := Decrypt(encryptedContent, tt.key)
+
+			decryptedContent, err := xor.Decrypt(encryptedContent, tt.key)
+			require.NoError(t, err)
 			require.Equal(t, tt.content, decryptedContent)
 		})
 	}

loader/aes.go

@@ -0,0 +1,41 @@
+package loader
+
+import (
+	"bytes"
+	"fmt"
+	"strings"
+	"text/template"
+
+	"github.com/cmepw/221b/encryption"
+	"github.com/cmepw/221b/templates"
+)
+
+type Aes struct {
+	baseLoader
+	encryption.Aes
+}
+
+func (a Aes) Load(content, key []byte) ([]byte, error) {
+	tmpl, err := template.New("loader").Funcs(template.FuncMap{
+		"key": func() string {
+			return string(key)
+		},
+		"shellcode": func() string {
+			result := []string{}
+
+			for _, b := range content {
+				result = append(result, fmt.Sprintf("0x%02x", b))
+			}
+
+			return strings.Join(result, ", ") + ","
+		},
+	}).Parse(templates.AesTmpl)
+	if err != nil {
+		return nil, err
+	}
+
+	result := new(bytes.Buffer)
+	err = tmpl.Execute(result, nil)
+
+	return result.Bytes(), err
+}

loader/loader.go

@@ -6,13 +6,21 @@ import (
 	"os/exec"
 	"path/filepath"
 
+	"github.com/cmepw/221b/encryption"
 	"github.com/cmepw/221b/logger"
 	"github.com/cmepw/221b/templates"
 )
 
+// Method gather all loading and encryption method.
+var Method = map[string]Loader{
+	"xor": Xor{},
+	"aes": Aes{},
+}
+
 type Loader interface {
-	Load(content []byte) ([]byte, error)
+	Load(content, key []byte) ([]byte, error)
 	Compile(path string, content []byte) error
+	encryption.Encryption
 }
 
 type baseLoader struct{}

loader/xor.go

@@ -6,22 +6,19 @@ import (
 	"strings"
 	"text/template"
 
+	"github.com/cmepw/221b/encryption"
 	"github.com/cmepw/221b/templates"
 )
 
 type Xor struct {
-	key []byte
 	baseLoader
+	encryption.Xor
 }
 
-func NewXorLoader(key []byte) *Xor {
-	return &Xor{key: key}
-}
-
-func (x Xor) Load(content []byte) ([]byte, error) {
+func (x Xor) Load(content, key []byte) ([]byte, error) {
 	tmpl, err := template.New("loader").Funcs(template.FuncMap{
 		"key": func() string {
-			return string(x.key)
+			return string(key)
 		},
 		"shellcode": func() string {
 			result := []string{}