PE-Injector/loader_test.go at main · RTS-Framework/PE-Injector · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
package injector

import (
	"fmt"
	"os"
	"testing"

	"github.com/stretchr/testify/require"
)

func TestLoader(t *testing.T) {
	injector := NewInjector()

	opts := &Options{
		NotSaveContext: true,
		NoGarbageInst:  true,
		RandSeed:       1234,
	}

	t.Run("auto", func(t *testing.T) {
		opts.ForceCodeCave = false
		opts.ForceCodeCaveNS = false
		opts.ForceExtendText = false
		opts.ForceExtendTextNS = false
		opts.ForceCreateText = false
		testLoader(t, injector, opts, "auto")
	})

	t.Run("code cave", func(t *testing.T) {
		opts.ForceCodeCave = true
		opts.ForceCodeCaveNS = false
		opts.ForceExtendText = false
		opts.ForceExtendTextNS = false
		opts.ForceCreateText = false
		testLoader(t, injector, opts, ModeCodeCave)
	})

	t.Run("code cave with new section", func(t *testing.T) {
		opts.ForceCodeCave = false
		opts.ForceCodeCaveNS = true
		opts.ForceExtendText = false
		opts.ForceExtendTextNS = false
		opts.ForceCreateText = false
		testLoader(t, injector, opts, ModeCodeCaveNS)
	})

	t.Run("extend text section", func(t *testing.T) {
		opts.ForceCodeCave = false
		opts.ForceCodeCaveNS = false
		opts.ForceExtendText = true
		opts.ForceExtendTextNS = false
		opts.ForceCreateText = false
		testLoader(t, injector, opts, ModeExtendText)
	})

	t.Run("extend text with new section", func(t *testing.T) {
		opts.ForceCodeCave = false
		opts.ForceCodeCaveNS = false
		opts.ForceExtendText = false
		opts.ForceExtendTextNS = true
		opts.ForceCreateText = false
		testLoader(t, injector, opts, ModeExtendTextNS)
	})

	t.Run("create text section", func(t *testing.T) {
		opts.ForceCodeCave = false
		opts.ForceCodeCaveNS = false
		opts.ForceExtendText = false
		opts.ForceExtendTextNS = false
		opts.ForceCreateText = true
		testLoader(t, injector, opts, ModeCreateText)
	})

	err := injector.Close()
	require.NoError(t, err)
}

func testLoader(t *testing.T, injector *Injector, opts *Options, mode string) {
	t.Run("x86", func(t *testing.T) {
		if opts.ForceCodeCave {
			return
		}

		expected := mode
		if expected == "auto" {
			expected = ModeCodeCaveNS
		}

		image, err := os.ReadFile("testdata/image_exe_x86.dat")
		require.NoError(t, err)
		shellcode, err := os.ReadFile("testdata/shellcode_x86.dat")
		require.NoError(t, err)

		ctx, err := injector.Inject(image, shellcode, opts)
		require.NoError(t, err)
		require.Equal(t, expected, ctx.Mode)
		fmt.Println(ctx.LoaderHex)
		fmt.Println(ctx.LoaderInst)

		testExecuteEXE(t, "testdata/injected_x86.exe", ctx.Output)
		testCheckOutputImage(t, image, ctx.Output, ctx.Mode)
	})

	t.Run("x64", func(t *testing.T) {
		expected := mode
		if expected == "auto" {
			expected = ModeCodeCave
		}

		image, err := os.ReadFile("testdata/image_exe_x64.dat")
		require.NoError(t, err)
		shellcode, err := os.ReadFile("testdata/shellcode_x64.dat")
		require.NoError(t, err)

		ctx, err := injector.Inject(image, shellcode, opts)
		require.NoError(t, err)
		require.Equal(t, expected, ctx.Mode)
		fmt.Println(ctx.LoaderHex)
		fmt.Println(ctx.LoaderInst)

		testExecuteEXE(t, "testdata/injected_x64.exe", ctx.Output)
		testCheckOutputImage(t, image, ctx.Output, ctx.Mode)
	})
}