Skip to content

Add Unit Testing for builder.go #534

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: dev
Choose a base branch
from
Open

Add Unit Testing for builder.go #534

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions teamserver/go.mod
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ require (
github.com/olekukonko/tablewriter v0.0.5
github.com/sergi/go-diff v1.2.0
github.com/spf13/cobra v1.2.1
github.com/stretchr/testify v1.7.0
github.com/zclconf/go-cty v1.9.0
golang.org/x/crypto v0.0.0-20220314234659-1baeb1ce4c0b
golang.org/x/image v0.5.0
Expand All @@ -39,9 +40,12 @@ require (
github.com/mattn/go-runewidth v0.0.9 // indirect
github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421 // indirect
github.com/modern-go/reflect2 v1.0.1 // indirect
github.com/pmezard/go-difflib v1.0.0 // indirect
github.com/spf13/pflag v1.0.5 // indirect
github.com/stretchr/objx v0.1.0 // indirect
github.com/ugorji/go/codec v1.1.7 // indirect
golang.org/x/sys v0.5.0 // indirect
google.golang.org/protobuf v1.26.0 // indirect
gopkg.in/yaml.v2 v2.4.0 // indirect
gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b // indirect
)
1 change: 1 addition & 0 deletions teamserver/go.sum
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -264,6 +264,7 @@ github.com/spf13/jwalterweatherman v1.1.0/go.mod h1:aNWZUN0dPAAO/Ljvb5BEdw96iTZ0
github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
github.com/spf13/viper v1.8.1/go.mod h1:o0Pch8wJ9BVSWGQMbra6iw0oQ5oktSIBaujf1rJH9Ns=
github.com/stretchr/objx v0.1.0 h1:4G4v2dO3VZwixGIRoQ5Lfboy6nUhCyYzaqnIAPPhYs4=
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
Expand Down
62 changes: 51 additions & 11 deletions teamserver/pkg/common/builder/builder.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@ package builder

import (
"bytes"
"io/fs"

//"encoding/hex"
"encoding/json"
"errors"
Expand Down Expand Up @@ -67,6 +69,32 @@ const (
ARCHITECTURE_X86 = 2
)

// Test Injection Interfaces

type OSOperations interface {
Mkdir(path string, perm os.FileMode) error
ReadDir(name string) ([]fs.DirEntry, error)
ReadFile(filename string) ([]byte, error)
Remove(name string) error
}

type FilePathOperations interface {
Abs(path string) (string, error)
}

// Default Implementations of Test Injections

type RealOSOperations struct{}

func (RealOSOperations) Mkdir(path string, perm os.FileMode) error { return os.Mkdir(path, perm) }
func (RealOSOperations) ReadDir(name string) ([]fs.DirEntry, error) { return os.ReadDir(name) }
func (RealOSOperations) ReadFile(filename string) ([]byte, error) { return os.ReadFile(filename) }
func (RealOSOperations) Remove(name string) error { return os.Remove(name) }

type RealFilePathOperations struct{}

func (RealFilePathOperations) Abs(path string) (string, error) { return filepath.Abs(path) }

type BuilderConfig struct {
Compiler64 string
Compiler86 string
Expand All @@ -76,6 +104,12 @@ type BuilderConfig struct {
}

type Builder struct {
osOps OSOperations
fpOps FilePathOperations
PatchConfigImpl func() ([]byte, error)
CompileCmdImpl func(cmd string) bool
CmdImpl func(cmd string) bool

buildSource bool
sourcePath string
silent bool
Expand Down Expand Up @@ -141,6 +175,12 @@ type Builder struct {
func NewBuilder(config BuilderConfig) *Builder {
var builder = new(Builder)

builder.osOps = &RealOSOperations{}
builder.fpOps = &RealFilePathOperations{}
builder.PatchConfigImpl = builder.PatchConfig
builder.CompileCmdImpl = builder.CompileCmd
builder.CmdImpl = builder.Cmd

builder.sourcePath = utils.GetTeamserverPath() + "/" + PayloadDir + "/Demon"
builder.config.Arch = ARCHITECTURE_X64

Expand Down Expand Up @@ -221,7 +261,7 @@ func (b *Builder) Build() bool {
)

b.CompileDir = "/tmp/" + utils.GenerateID(10) + "/"
err := os.Mkdir(b.CompileDir, os.ModePerm)
err := b.osOps.Mkdir(b.CompileDir, os.ModePerm)
if err != nil {
logger.Error("Failed to create compile directory: " + err.Error())
return false
Expand All @@ -241,7 +281,7 @@ func (b *Builder) Build() bool {
b.SendConsoleMessage("Info", "starting build")
}

Config, err := b.PatchConfig()
Config, err := b.PatchConfigImpl()
if err != nil {
b.SendConsoleMessage("Error", err.Error())
return false
Expand Down Expand Up @@ -283,7 +323,7 @@ func (b *Builder) Build() bool {

// add compiler
if b.config.Arch == ARCHITECTURE_X64 {
abs, err := filepath.Abs(b.compilerOptions.Config.Compiler64)
abs, err := b.fpOps.Abs(b.compilerOptions.Config.Compiler64)

if err != nil {
if !b.silent {
Expand All @@ -295,7 +335,7 @@ func (b *Builder) Build() bool {

CompileCommand += "\"" + b.compilerOptions.Config.Compiler64 + "\" "
} else {
abs, err := filepath.Abs(b.compilerOptions.Config.Compiler86)
abs, err := b.fpOps.Abs(b.compilerOptions.Config.Compiler86)

if err != nil {
if !b.silent {
Expand All @@ -310,7 +350,7 @@ func (b *Builder) Build() bool {

// add sources
for _, dir := range b.compilerOptions.SourceDirs {
files, err := os.ReadDir(b.sourcePath + "/" + dir)
files, err := b.osOps.ReadDir(b.sourcePath + "/" + dir)
if err != nil {
logger.Error(err)
}
Expand All @@ -330,7 +370,7 @@ func (b *Builder) Build() bool {
}
logger.Debug(AsmCompile)
b.FilesCreated = append(b.FilesCreated, AsmObj)
b.Cmd(AsmCompile)
b.CmdImpl(AsmCompile)
CompileCommand += AsmObj + " "
}
} else if path.Ext(f.Name()) == ".c" {
Expand Down Expand Up @@ -426,7 +466,7 @@ func (b *Builder) Build() bool {
ShellcodePath = utils.GetTeamserverPath() + "/" + PayloadDir + "/Shellcode.x86.bin"
}

ShellcodeTemplate, err := os.ReadFile(ShellcodePath)
ShellcodeTemplate, err := b.osOps.ReadFile(ShellcodePath)
if err != nil {
logger.Error("Couldn't read content of file: " + err.Error())
b.SendConsoleMessage("Error", "couldn't read content of file: "+err.Error())
Expand All @@ -451,7 +491,7 @@ func (b *Builder) Build() bool {
}

//logger.Debug(CompileCommand)
Successful := b.CompileCmd(CompileCommand)
Successful := b.CompileCmdImpl(CompileCommand)

return Successful
}
Expand Down Expand Up @@ -1041,7 +1081,7 @@ func (b *Builder) GetPayloadBytes() []byte {
return nil
}

FileBuffer, err = os.ReadFile(b.outputPath)
FileBuffer, err = b.osOps.ReadFile(b.outputPath)
if err != nil {
logger.Error("Couldn't read content of file: " + err.Error())
if !b.silent {
Expand Down Expand Up @@ -1089,7 +1129,7 @@ func (b *Builder) Cmd(cmd string) bool {

func (b *Builder) CompileCmd(cmd string) bool {

if b.Cmd(cmd) {
if b.CmdImpl(cmd) {
if !b.silent {
b.SendConsoleMessage("Info", "finished compiling source")
}
Expand Down Expand Up @@ -1124,7 +1164,7 @@ func (b *Builder) DeletePayload() {
b.FilesCreated = append(b.FilesCreated, b.CompileDir)
for _, FileCreated := range b.FilesCreated {
if strings.HasSuffix(FileCreated, ".bin") == false {
if err := os.Remove(FileCreated); err != nil {
if err := b.osOps.Remove(FileCreated); err != nil {
logger.Debug("Couldn't remove " + FileCreated + ": " + err.Error())
}
}
Expand Down
Loading