admin/FalconOpsLLC-goexec
1
0
Fork 0
mirror of https://github.com/FalconOpsLLC/goexec synced 2026-06-06 15:44:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add hidden cpu & memory profiling flags

Browse Source
This commit is contained in:
Bryan McNulty
2025-04-25 18:58:06 -05:00
parent 3c6d9d2244
commit 69564ac0f9
1 changed files with 64 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cmd/root.go
+64 -3
View File
@@ -15,6 +15,7 @@ import (
"golang.org/x/term"
"io"
"os"
"runtime/pprof"
)
type flagSet struct {
@@ -56,6 +57,7 @@ var (
defaultAuthFlags, defaultLogFlags, defaultNetRpcFlags *flagSet
returnCode int
toClose []io.Closer
// === IO ===
stageFilePath string
@@ -79,6 +81,13 @@ var (
smbClient smb.Client
// ===============
// === Resource profiling ===
cpuProfile string
memProfile string
cpuProfileFile io.WriteCloser
memProfileFile io.WriteCloser
// ==========================
exec = goexec.ExecutionIO{
Input: new(goexec.ExecutionInput),
Output: new(goexec.ExecutionOutput),
@@ -113,6 +122,7 @@ Authors: FalconOps LLC (@FalconOpsLLC),
if err != nil {
return
}
toClose = append(toClose, logFile)
logJson = true
}
if logQuiet {
@@ -128,6 +138,29 @@ Authors: FalconOps LLC (@FalconOpsLLC),
log = log.Level(logLevel)
}
// CPU / memory profiling
{
if cpuProfile != "" {
if cpuProfileFile, err = os.OpenFile(cpuProfile, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644); err != nil {
log.Error().Err(err).Msg("Failed to open CPU profile for writing")
return
}
toClose = append(toClose, cpuProfileFile)
if err = pprof.StartCPUProfile(cpuProfileFile); err != nil {
log.Error().Err(err).Msg("Failed to start CPU profile")
return
}
}
if memProfile != "" {
if memProfileFile, err = os.OpenFile(memProfile, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644); err != nil {
log.Error().Err(err).Msg("Failed to open memory profile for writing")
return
}
toClose = append(toClose, memProfileFile)
}
}
if proxy != "" {
rpcClient.Proxy = proxy
smbClient.Proxy = proxy
@@ -151,11 +184,26 @@ Authors: FalconOps LLC (@FalconOpsLLC),
},
PersistentPostRun: func(cmd *cobra.Command, args []string) {
if logFile != nil {
if err := logFile.Close(); err != nil {
// ...
if memProfileFile != nil {
if err := pprof.WriteHeapProfile(memProfileFile); err != nil {
log.Error().Err(err).Msg("Failed to write memory profile")
return
}
}
if cpuProfileFile != nil {
pprof.StopCPUProfile()
}
for _, c := range toClose {
if c != nil {
if err := c.Close(); err != nil {
log.Warn().Err(err).Msg("Failed to close stream")
}
}
}
if exec.Input != nil && exec.Input.StageFile != nil {
if err := exec.Input.StageFile.Close(); err != nil {
// ...
@@ -182,6 +230,19 @@ func init() {
gssapi.AddMechanism(ssp.KRB5)
}
// CPU / Memory profiling
{
rootCmd.PersistentFlags().StringVar(&cpuProfile, "cpu-profile", "", "Write CPU profile to `file`")
rootCmd.PersistentFlags().StringVar(&memProfile, "mem-profile", "", "Write memory profile to `file`")
if err := rootCmd.PersistentFlags().MarkHidden("cpu-profile"); err != nil {
panic(err)
}
if err := rootCmd.PersistentFlags().MarkHidden("mem-profile"); err != nil {
panic(err)
}
}
// Cobra init
{
cobra.EnableCommandSorting = false