chore: update gonut check pe type

Author: M09Ic

external/gonut/donut.go

@@ -2,27 +2,17 @@ package gonut
 
 import (
 	"os"
-	"path/filepath"
 	"strings"
 )
 
 // DonutShellcodeFromFile 从给定的 PE 文件生成 Donut shellcode
 func DonutShellcodeFromFile(filePath string, arch string, params string) (data []byte, err error) {
-	pe, err := os.ReadFile(filePath)
+	bin, err := os.ReadFile(filePath)
 	if err != nil {
 		return
 	}
-	ext, err := GetExtension(filePath)
-	if err != nil {
-		return
-	}
-
-	baseName := filepath.Base(filePath)
-	if filepath.Ext(baseName) == "" && ext != "" {
-		baseName = baseName + ext
-	}
 
-	return DonutShellcodeFromPE(baseName, pe, arch, params, false, true)
+	return DonutShellcodeFromPE(filePath, bin, arch, params, false, true)
 }
 
 // DonutShellcodeFromPE 从给定的 PE 数据生成 Donut shellcode

external/gonut/utils.go

@@ -71,7 +71,6 @@ func UnsafeStructToBytes(ptr any) []byte {
 }
 
 func GetExtension(filepath string) (string, error) {
-
 	file, err := os.Open(filepath)
 	if err != nil {
 		return "", fmt.Errorf("failed to open file %s: %w", filepath, err)

helper/consts/implant.go

@@ -126,7 +126,6 @@ const (
 
 const (
 	ELF           = ".elf"
-	PE            = ".pe"
 	DLL           = ".dll"
 	PEFile        = ".exe"
 	ShellcodeFile = ".bin"

helper/utils/formatutils/formatter.go

@@ -7,8 +7,6 @@ import (
 	"github.com/chainreactors/malice-network/helper/cryptography"
 	"github.com/chainreactors/malice-network/helper/encoders"
 	"github.com/chainreactors/malice-network/helper/proto/client/clientpb"
-	"os"
-	"path/filepath"
 	"strings"
 )
 
@@ -156,17 +154,12 @@ func ConvertArtifact(artifact *clientpb.Artifact, format string) (*clientpb.Arti
 		artifact.Format = format
 		return artifact, nil
 	}
-	filename := filepath.Join(encoders.UUID())
-	if err := os.WriteFile(filename, artifact.Bin, 0644); err != nil {
-		return nil, err
-	}
-	shellcode, err := SRDIArtifact(filename, artifact.Platform, artifact.Arch, artifact.Type == consts.CommandBuildPulse)
+
+	shellcode, err := SRDIArtifact(artifact.Bin, artifact.Platform, artifact.Arch, artifact.Type == consts.CommandBuildPulse)
 	if err != nil {
 		return nil, fmt.Errorf("failed to convert: %s", err)
 	}
-	if err := os.Remove(filename); err != nil {
-		return nil, fmt.Errorf("failed to remove file: %s", err)
-	}
+
 	convert, err := Convert(shellcode, format)
 	if err != nil {
 		return nil, err

helper/utils/formatutils/srdi.go

@@ -3,6 +3,9 @@ package formatutils
 import (
 	"fmt"
 	"github.com/chainreactors/logs"
+	"github.com/chainreactors/malice-network/helper/consts"
+	"github.com/chainreactors/malice-network/helper/encoders"
+	"github.com/chainreactors/malice-network/helper/utils/pe"
 	"github.com/wabzsy/gonut"
 	"os"
 	"os/exec"
@@ -52,10 +55,22 @@ func ObjcopyPulse(path, platform, arch string) ([]byte, error) {
 	return bin, nil
 }
 
-func SRDIArtifact(path, platform, arch string, useobjcopy bool) ([]byte, error) {
+func SRDIArtifact(bin []byte, platform, arch string, useobjcopy bool) ([]byte, error) {
 	if useobjcopy {
-		return ObjcopyPulse(path, platform, arch)
+		filename := filepath.Join(encoders.UUID())
+		if err := os.WriteFile(filename, bin, 0644); err != nil {
+			return nil, err
+		}
+		defer os.Remove(filename)
+		return ObjcopyPulse(filename, platform, arch)
 	} else {
-		return gonut.DonutShellcodeFromFile(path, arch, "")
+		switch pe.CheckPEType(bin) {
+		case consts.DLLFile:
+			return gonut.DonutShellcodeFromPE("bin"+consts.DLL, bin, arch, "", false, true)
+		case consts.EXEFile:
+			return gonut.DonutShellcodeFromPE("bin"+consts.PEFile, bin, arch, "", false, true)
+		default:
+			return nil, fmt.Errorf("unsupported file type")
+		}
 	}
 }