diff --git a/cmd/pipeleek/main.go b/cmd/pipeleek/main.go index 4db8e438..ca159957 100644 --- a/cmd/pipeleek/main.go +++ b/cmd/pipeleek/main.go @@ -1,64 +1,16 @@ package main import ( - "encoding/json" - "io" - "os" - "github.com/CompassSecurity/pipeleek/internal/cmd" - "golang.org/x/term" + "github.com/CompassSecurity/pipeleek/pkg/logging" ) -var originalTermState *term.State - -type TerminalRestoringWriter struct { - underlying io.Writer -} - -func (w *TerminalRestoringWriter) Write(p []byte) (n int, err error) { - var logEntry map[string]interface{} - if err := json.Unmarshal(p, &logEntry); err == nil { - if level, ok := logEntry["level"].(string); ok && level == "fatal" { - _, _ = w.underlying.Write(p) - restoreTerminalState() - os.Exit(1) - } - } - return w.underlying.Write(p) -} - func main() { - saveTerminalState() - defer restoreTerminalState() - - cmd.TerminalRestorer = restoreTerminalState + logging.InstallExitHandler() + defer logging.StopShortcutListeners() err := cmd.Execute() if err != nil { - os.Exit(1) + logging.Exit(1) } } - -func saveTerminalState() { - // #nosec G115 -- os.Stdin file descriptor is provided by the runtime and fits into int on supported platforms - stdinFD := int(os.Stdin.Fd()) - if term.IsTerminal(stdinFD) { - state, err := term.GetState(stdinFD) - if err == nil { - originalTermState = state - } - } -} - -func restoreTerminalState() { - if originalTermState != nil { - // #nosec G115 -- os.Stdin file descriptor is provided by the runtime and fits into int on supported platforms - stdinFD := int(os.Stdin.Fd()) - _ = term.Restore(stdinFD, originalTermState) - } -} - -// GetTerminalRestoringWriter wraps the given writer to restore terminal state on fatal logs -func GetTerminalRestoringWriter(w io.Writer) io.Writer { - return &TerminalRestoringWriter{underlying: w} -} diff --git a/docs/introduction/logging.md b/docs/introduction/logging.md index 3fadd9e5..37df1cbf 100644 --- a/docs/introduction/logging.md +++ b/docs/introduction/logging.md @@ -91,6 +91,11 @@ pipeleek gl scan -u https://gitlab.com -t glpat-xxxxx --log-level=trace You can change interactively between log levels by pressing `t`: Trace, `d`: Debug, `i`: info, `w`: Warn, `e`: Error. +Keyboard shortcuts are enabled only when standard input is a terminal. Pressing +Escape or Ctrl+C stops the shortcut listener and restores normal terminal input. +Terminal input is also restored when a command finishes, returns an error, logs +a fatal error, or completes graceful shutdown. + ```bash pipeleek gl scan -u https://gitlab.com -t glpat-[redacted] --secrets-verification=false --verbose # Human Pressed d on keyboard diff --git a/go.mod b/go.mod index c5aa3171..1f81cdd5 100644 --- a/go.mod +++ b/go.mod @@ -3,7 +3,6 @@ module github.com/CompassSecurity/pipeleek go 1.26.0 require ( - atomicgo.dev/keyboard v0.2.10 code.gitea.io/sdk/gitea v0.25.1 github.com/CircleCI-Public/circleci-cli v0.1.38646 github.com/PuerkitoBio/goquery v1.13.0 @@ -76,7 +75,7 @@ require ( github.com/charlievieth/fastwalk v1.0.14 // indirect github.com/clipperhouse/uax29/v2 v2.7.0 // indirect github.com/cloudflare/circl v1.6.3 // indirect - github.com/containerd/console v1.0.5 // indirect + github.com/containerd/console v1.0.5 github.com/couchbase/gocb/v2 v2.11.0 // indirect github.com/couchbase/gocbcore/v10 v10.8.0 // indirect github.com/couchbase/gocbcoreps v0.1.3 // indirect @@ -237,7 +236,7 @@ require ( github.com/wandb/parallel v0.2.3 github.com/yosuke-furukawa/json5 v0.1.1 golang.org/x/oauth2 v0.36.0 // indirect - golang.org/x/sys v0.49.0 // indirect + golang.org/x/sys v0.49.0 golang.org/x/time v0.14.0 // indirect google.golang.org/protobuf v1.36.11 // indirect gopkg.in/yaml.v3 v3.0.1 diff --git a/go.sum b/go.sum index 77501b23..c6c2405d 100644 --- a/go.sum +++ b/go.sum @@ -1,11 +1,3 @@ -atomicgo.dev/assert v0.0.2 h1:FiKeMiZSgRrZsPo9qn/7vmr7mCsh5SZyXY4YGYiYwrg= -atomicgo.dev/assert v0.0.2/go.mod h1:ut4NcI3QDdJtlmAxQULOmA13Gz6e2DWbSAS8RUOmNYQ= -atomicgo.dev/cursor v0.2.0 h1:H6XN5alUJ52FZZUkI7AlJbUc1aW38GWZalpYRPpoPOw= -atomicgo.dev/cursor v0.2.0/go.mod h1:Lr4ZJB3U7DfPPOkbH7/6TOtJ4vFGHlgj1nc+n900IpU= -atomicgo.dev/keyboard v0.2.10 h1:v7mvUKUZLHIggxULEIuWbT+WkkyQSgdbA201EziAhHU= -atomicgo.dev/keyboard v0.2.10/go.mod h1:ap/z5ilnhLqYq852m6kPeTq5Z6aESGWu5mzRpJlC6aI= -atomicgo.dev/schedule v0.1.0 h1:nTthAbhZS5YZmgYbb2+DH8uQIZcTlIrd4eYr3UQxEjs= -atomicgo.dev/schedule v0.1.0/go.mod h1:xeUa3oAkiuHYh8bKiQBRojqAMq3PXXbJujjb0hw8pEU= cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw= cloud.google.com/go v0.121.6 h1:waZiuajrI28iAf40cWgycWNgaXPO06dupuS+sgibK6c= cloud.google.com/go v0.121.6/go.mod h1:coChdst4Ea5vUpiALcYKXEpR1S9ZgXbhEzzMcMR66vI= @@ -48,8 +40,6 @@ github.com/BobuSumisu/aho-corasick v1.0.3/go.mod h1:hm4jLcvZKI2vRF2WDU1N4p/jpWtp github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU= github.com/CircleCI-Public/circleci-cli v0.1.38646 h1:kymjys07c2Vk3y/BnHpMz24sU0zhpGlwJm5n76KdJrs= github.com/CircleCI-Public/circleci-cli v0.1.38646/go.mod h1:XeQQRIYyT7CYNMCI3BAkURiUxgzZZ1EKWqPkO9pTHNs= -github.com/MarvinJWendt/testza v0.5.2 h1:53KDo64C1z/h/d/stCYCPY69bt/OSwjq5KpFNwi+zB4= -github.com/MarvinJWendt/testza v0.5.2/go.mod h1:xu53QFE5sCdjtMCKk8YMQ2MnymimEctc4n3EjyIYvEY= github.com/Microsoft/go-winio v0.5.2/go.mod h1:WpS1mjBmmwHBEWmogvA2mj8546UReBk4v8QkMxJ6pZY= github.com/Microsoft/go-winio v0.6.2 h1:F2VQgta7ecxGYO8k3ZZz3RS8fVIXVxONVUPlNERoyfY= github.com/Microsoft/go-winio v0.6.2/go.mod h1:yd8OoFMLzJbo9gZq8j5qaps8bJ9aShtEA8Ipt1oGCvU= @@ -296,8 +286,6 @@ github.com/googleapis/enterprise-certificate-proxy v0.3.11 h1:vAe81Msw+8tKUxi2Dq github.com/googleapis/enterprise-certificate-proxy v0.3.11/go.mod h1:RFV7MUdlb7AgEq2v7FmMCfeSMCllAzWxFgRdusoGks8= github.com/googleapis/gax-go/v2 v2.17.0 h1:RksgfBpxqff0EZkDWYuz9q/uWsTVz+kf43LsZ1J6SMc= github.com/googleapis/gax-go/v2 v2.17.0/go.mod h1:mzaqghpQp4JDh3HvADwrat+6M3MOIDp5YKHhb9PAgDY= -github.com/gookit/color v1.6.0 h1:JjJXBTk1ETNyqyilJhkTXJYYigHG24TM9Xa2M1xAhRA= -github.com/gookit/color v1.6.0/go.mod h1:9ACFc7/1IpHGBW8RwuDm/0YEnhg3dwwXpoMsmtyHfjs= github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg= github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= github.com/graph-gophers/graphql-go v1.9.0 h1:yu0ucKHLc5qGpRwLYKIWtr9bOoxovkWasuBrPQwlHls= @@ -356,10 +344,7 @@ github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+o github.com/klauspost/compress v1.4.1/go.mod h1:RyIbtBH6LamlWaDj8nUwkbUhJ87Yi3uG0guNDohfE1A= github.com/klauspost/compress v1.20.0 h1:a3C1ke2ohxFymNlb2HWAHjDeKCI90scRskErZkR0ezA= github.com/klauspost/compress v1.20.0/go.mod h1:LUdAzn7YLVvxLpc7y3V1m40wESHTgc1422pwwBSKYuI= -github.com/klauspost/cpuid v1.2.0 h1:NMpwD2G9JSFOE1/TJjGSo5zG7Yb2bTe7eq1jH+irmeE= github.com/klauspost/cpuid v1.2.0/go.mod h1:Pj4uuM528wm8OyEC2QMXAi2YiTZ96dNQPGgoMS4s3ek= -github.com/klauspost/cpuid/v2 v2.3.0 h1:S4CRMLnYUhGeDFDqkGriYKdfoFlDnMtqTiI/sFzhA9Y= -github.com/klauspost/cpuid/v2 v2.3.0/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0= github.com/klauspost/pgzip v1.2.6 h1:8RXeL5crjEUFnR2/Sn6GJNWtSQ3Dk8pq4CL3jvdDyjU= github.com/klauspost/pgzip v1.2.6/go.mod h1:Ch1tH69qFZu15pkjo5kYi6mth2Zzwzt50oCQKQE9RUs= github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ= @@ -390,8 +375,6 @@ github.com/lestrrat-go/option/v2 v2.0.0 h1:XxrcaJESE1fokHy3FpaQ/cXW8ZsIdWcdFzzLO github.com/lestrrat-go/option/v2 v2.0.0/go.mod h1:oSySsmzMoR0iRzCDCaUfsCzxQHUEuhOViQObyy7S6Vg= github.com/lib/pq v1.10.9 h1:YXG7RB+JIjhP29X+OtkiDnYaXQwpS4JEWq7dtCCRUEw= github.com/lib/pq v1.10.9/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o= -github.com/lithammer/fuzzysearch v1.1.8 h1:/HIuJnjHuXS8bKaiTMeeDlW2/AyIWk2brx1V8LFgLN4= -github.com/lithammer/fuzzysearch v1.1.8/go.mod h1:IdqeyBClc3FFqSzYq/MXESsS4S0FsZ5ajtkr5xPLts4= github.com/lucasjones/reggen v0.0.0-20200904144131-37ba4fa293bb h1:w1g9wNDIE/pHSTmAaUhv4TZQuPBS6GV3mMz5hkgziIU= github.com/lucasjones/reggen v0.0.0-20200904144131-37ba4fa293bb/go.mod h1:5ELEyG+X8f+meRWHuqUOewBOhvHkl7M76pdGEansxW4= github.com/lufia/plan9stats v0.0.0-20260330125221-c963978e514e h1:Q6MvJtQK/iRcRtzAscm/zF23XxJlbECiGPyRicsX+Ak= @@ -494,8 +477,6 @@ github.com/prometheus/common v0.62.0 h1:xasJaQlnWAeyHdUBeGjXmutelfJHWMRr+Fg4QszZ github.com/prometheus/common v0.62.0/go.mod h1:vyBcEuLSvWos9B1+CyL7JZ2up+uFzXhkqml0W5zIY1I= github.com/prometheus/procfs v0.15.1 h1:YagwOFzUgYfKKHX6Dr+sHT7km/hxC76UB0learggepc= github.com/prometheus/procfs v0.15.1/go.mod h1:fB45yRUv8NstnjriLhBQLuOUt+WW4BsoGhij/e3PBqk= -github.com/pterm/pterm v0.12.83 h1:ie+YmGmA727VuhxBlyGr74Ks+7McV6kT99IB8EU80aA= -github.com/pterm/pterm v0.12.83/go.mod h1:xlgc6bFWyJIMtmLJvGim+L7jhSReilOlOnodeIYe4Tk= github.com/rabbitmq/amqp091-go v1.14.0 h1:RSaT7aOKt/OrkVUyswPDW29lnRz9psuGmfZFBmLqLek= github.com/rabbitmq/amqp091-go v1.14.0/go.mod h1:Hy4jKW5kQART1u+JkDTF9YYOQUHXqMuhrgxOEeS7G4o= github.com/rhysd/actionlint v1.7.12 h1:vQ4GeJN86C0QH+gTUQcs8McmK62OLT3kmakPMtEWYnY= @@ -618,8 +599,6 @@ github.com/xdg-go/scram v1.1.2 h1:FHX5I5B4i4hKRVRBCFRxq1iQRej7WO3hhBuJf+UUySY= github.com/xdg-go/scram v1.1.2/go.mod h1:RT/sEzTbU5y00aCK8UOx6R7YryM0iF1N2MOmC3kKLN4= github.com/xdg-go/stringprep v1.0.4 h1:XLI/Ng3O1Atzq0oBs3TWm+5ZVgkq2aqdlvP9JtoZ6c8= github.com/xdg-go/stringprep v1.0.4/go.mod h1:mPGuuIYwz7CmR2bT9j4GbQqutWS1zV24gijq1dTyGkM= -github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no= -github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM= github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZqKjWU= github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E= github.com/yosuke-furukawa/json5 v0.1.1 h1:0F9mNwTvOuDNH243hoPqvf+dxa5QsKnZzU20uNsh3ZI= diff --git a/internal/cmd/common/common.go b/internal/cmd/common/common.go index aec2bd9d..e5e71fbf 100644 --- a/internal/cmd/common/common.go +++ b/internal/cmd/common/common.go @@ -3,7 +3,6 @@ package common import ( "bytes" - "encoding/json" "io" "os" "runtime" @@ -15,7 +14,6 @@ import ( "github.com/rs/zerolog" "github.com/rs/zerolog/log" "github.com/spf13/cobra" - "golang.org/x/term" ) // Version information - set via ldflags during build @@ -27,20 +25,16 @@ var ( // Log configuration var ( - originalTermState *term.State - JsonLogoutput bool - LogFile string - LogColor bool - LogDebug bool - LogLevel string - IgnoreProxy bool + JsonLogoutput bool + LogFile string + LogColor bool + LogDebug bool + LogLevel string + IgnoreProxy bool // logFileCloser holds the currently open log file so it can be closed on cleanup. logFileCloser io.Closer ) -// TerminalRestorer is a function that can be called to restore terminal state -var TerminalRestorer func() - // CustomWriter wraps an os.File with proper cross-platform newline handling type CustomWriter struct { Writer *os.File @@ -73,55 +67,6 @@ func (cw *CustomWriter) Write(p []byte) (n int, err error) { return originalLen, nil } -// TerminalRestoringWriter wraps an io.Writer to restore terminal state on fatal logs -type TerminalRestoringWriter struct { - underlying io.Writer -} - -func (w *TerminalRestoringWriter) Write(p []byte) (n int, err error) { - var logEntry map[string]interface{} - if err := json.Unmarshal(p, &logEntry); err == nil { - if level, ok := logEntry["level"].(string); ok && level == "fatal" { - _, _ = w.underlying.Write(p) - RestoreTerminalState() - os.Exit(1) - } - } - return w.underlying.Write(p) -} - -// FatalHook is a zerolog hook that restores terminal state before fatal exits -type FatalHook struct{} - -func (h FatalHook) Run(e *zerolog.Event, level zerolog.Level, msg string) { - if level == zerolog.FatalLevel { - if TerminalRestorer != nil { - TerminalRestorer() - } - } -} - -// SaveTerminalState saves the current terminal state for later restoration -func SaveTerminalState() { - // #nosec G115 -- os.Stdin file descriptor is provided by the runtime and fits into int on supported platforms - stdinFD := int(os.Stdin.Fd()) - if term.IsTerminal(stdinFD) { - state, err := term.GetState(stdinFD) - if err == nil { - originalTermState = state - } - } -} - -// RestoreTerminalState restores the terminal to its saved state -func RestoreTerminalState() { - if originalTermState != nil { - // #nosec G115 -- os.Stdin file descriptor is provided by the runtime and fits into int on supported platforms - stdinFD := int(os.Stdin.Fd()) - _ = term.Restore(stdinFD, originalTermState) - } -} - // InitLogger initializes the zerolog logger with the configured options func InitLogger(cmd *cobra.Command) { defaultOut := &CustomWriter{Writer: os.Stdout} @@ -146,14 +91,12 @@ func InitLogger(cmd *cobra.Command) { } } - fatalHook := FatalHook{} - if JsonLogoutput { // For JSON output, wrap with HitLevelWriter to transform level field hitWriter := &logging.HitLevelWriter{} hitWriter.SetOutput(defaultOut) logging.SetGlobalHitWriter(hitWriter) - log.Logger = zerolog.New(hitWriter).With().Timestamp().Logger().Hook(fatalHook) + log.Logger = zerolog.New(hitWriter).With().Timestamp().Logger() } else { // For console output, use custom FormatLevel to color the hit level output := zerolog.ConsoleWriter{ @@ -166,7 +109,7 @@ func InitLogger(cmd *cobra.Command) { hitWriter := &logging.HitLevelWriter{} hitWriter.SetOutput(&output) logging.SetGlobalHitWriter(hitWriter) - log.Logger = zerolog.New(hitWriter).With().Timestamp().Logger().Hook(fatalHook) + log.Logger = zerolog.New(hitWriter).With().Timestamp().Logger() } } @@ -264,19 +207,17 @@ func SetupPersistentPreRun(cmd *cobra.Command) { InitLogger(c) SetGlobalLogLevel(c) httpclient.SetIgnoreProxy(IgnoreProxy) - go logging.ShortcutListeners(nil) + logging.ShortcutListeners(nil) } } // Run executes the common startup sequence and runs the provided root command func Run(rootCmd *cobra.Command) { - SaveTerminalState() - defer RestoreTerminalState() - - TerminalRestorer = RestoreTerminalState + logging.InstallExitHandler() + defer logging.StopShortcutListeners() err := rootCmd.Execute() if err != nil { - os.Exit(1) + logging.Exit(1) } } diff --git a/internal/cmd/common/common_test.go b/internal/cmd/common/common_test.go index 2ca33fb7..2f23a44a 100644 --- a/internal/cmd/common/common_test.go +++ b/internal/cmd/common/common_test.go @@ -1,7 +1,6 @@ package common import ( - "bytes" "os" "path/filepath" "testing" @@ -13,63 +12,6 @@ import ( "github.com/stretchr/testify/require" ) -func TestFatalHook_FatalLevelCallsTerminalRestorer(t *testing.T) { - orig := TerminalRestorer - defer func() { TerminalRestorer = orig }() - - called := false - TerminalRestorer = func() { called = true } - - hook := FatalHook{} - var buf bytes.Buffer - logger := zerolog.New(&buf) - event := logger.Fatal() - hook.Run(event, zerolog.FatalLevel, "test message") - - assert.True(t, called, "TerminalRestorer should be called on fatal level") -} - -func TestFatalHook_NonFatalLevelsDoNotCallTerminalRestorer(t *testing.T) { - orig := TerminalRestorer - defer func() { TerminalRestorer = orig }() - - levels := []zerolog.Level{ - zerolog.TraceLevel, - zerolog.DebugLevel, - zerolog.InfoLevel, - zerolog.WarnLevel, - zerolog.ErrorLevel, - } - - hook := FatalHook{} - var buf bytes.Buffer - logger := zerolog.New(&buf) - - for _, lvl := range levels { - called := false - TerminalRestorer = func() { called = true } - - event := logger.WithLevel(lvl) - hook.Run(event, lvl, "test") - assert.False(t, called, "TerminalRestorer should not be called for level %v", lvl) - } -} - -func TestFatalHook_NilTerminalRestorerDoesNotPanic(t *testing.T) { - orig := TerminalRestorer - defer func() { TerminalRestorer = orig }() - - TerminalRestorer = nil - hook := FatalHook{} - - assert.NotPanics(t, func() { - var buf bytes.Buffer - logger := zerolog.New(&buf) - event := logger.Fatal() - hook.Run(event, zerolog.FatalLevel, "test") - }) -} - func TestCustomWriter_WritesCorrectly(t *testing.T) { tmpDir := t.TempDir() logFile := filepath.Join(tmpDir, "test.log") @@ -393,29 +335,3 @@ func TestInitLogger_NoLogFile(t *testing.T) { InitLogger(cmd) }) } - -// TestSaveAndRestoreTerminalState verifies that SaveTerminalState and RestoreTerminalState -// do not panic regardless of whether stdin is a terminal. -func TestSaveAndRestoreTerminalState(t *testing.T) { - origState := originalTermState - defer func() { originalTermState = origState }() - - assert.NotPanics(t, func() { - SaveTerminalState() - }) - - assert.NotPanics(t, func() { - RestoreTerminalState() - }) -} - -// TestRestoreTerminalState_NilState verifies no panic when terminal state was never saved. -func TestRestoreTerminalState_NilState(t *testing.T) { - origState := originalTermState - defer func() { originalTermState = origState }() - - originalTermState = nil - assert.NotPanics(t, func() { - RestoreTerminalState() - }) -} diff --git a/internal/cmd/root.go b/internal/cmd/root.go index f3573cbb..f5c583b5 100644 --- a/internal/cmd/root.go +++ b/internal/cmd/root.go @@ -26,9 +26,6 @@ import ( "github.com/spf13/cobra" ) -// TerminalRestorer is a function that can be called to restore terminal state -var TerminalRestorer func() - // Version information - set via ldflags during build var ( Version = "dev" @@ -54,7 +51,7 @@ var ( httpclient.SetInsecureSkipVerify(!TLSVerification) httpclient.SetProxy(Proxy) httpclient.SetHTTPTimeout(HTTPTimeout) - go logging.ShortcutListeners(nil) + logging.ShortcutListeners(nil) }, } JsonLogoutput bool @@ -149,17 +146,6 @@ func (cw *CustomWriter) Write(p []byte) (n int, err error) { return originalLen, nil } -// FatalHook is a zerolog hook that restores terminal state before fatal exits -type FatalHook struct{} - -func (h FatalHook) Run(e *zerolog.Event, level zerolog.Level, msg string) { - if level == zerolog.FatalLevel { - if TerminalRestorer != nil { - TerminalRestorer() - } - } -} - func initLogger(cmd *cobra.Command) { defaultOut := &CustomWriter{Writer: os.Stdout} colorEnabled := LogColor @@ -186,14 +172,12 @@ func initLogger(cmd *cobra.Command) { // note: CloseLogger is defined at package scope - fatalHook := FatalHook{} - if JsonLogoutput { // For JSON output, wrap with HitLevelWriter to transform level field hitWriter := &logging.HitLevelWriter{} hitWriter.SetOutput(defaultOut) logging.SetGlobalHitWriter(hitWriter) - log.Logger = zerolog.New(hitWriter).With().Timestamp().Logger().Hook(fatalHook) + log.Logger = zerolog.New(hitWriter).With().Timestamp().Logger() } else { // For console output, use custom FormatLevel to color the hit level output := zerolog.ConsoleWriter{ @@ -206,7 +190,7 @@ func initLogger(cmd *cobra.Command) { hitWriter := &logging.HitLevelWriter{} hitWriter.SetOutput(&output) logging.SetGlobalHitWriter(hitWriter) - log.Logger = zerolog.New(hitWriter).With().Timestamp().Logger().Hook(fatalHook) + log.Logger = zerolog.New(hitWriter).With().Timestamp().Logger() } } diff --git a/internal/cmd/root_test.go b/internal/cmd/root_test.go index 5881a933..97887788 100644 --- a/internal/cmd/root_test.go +++ b/internal/cmd/root_test.go @@ -1,7 +1,6 @@ package cmd import ( - "bytes" "os" "path/filepath" "testing" @@ -139,92 +138,6 @@ func TestPersistentPreRunRegistered(t *testing.T) { assert.NotNil(t, rootCmd.PersistentPreRun, "PersistentPreRun should be registered") } -func TestTerminalRestorer(t *testing.T) { - t.Run("TerminalRestorer_can_be_set", func(t *testing.T) { - originalRestorer := TerminalRestorer - defer func() { TerminalRestorer = originalRestorer }() - - called := false - TerminalRestorer = func() { called = true } - - TerminalRestorer() - assert.True(t, called, "TerminalRestorer should be callable") - }) - - t.Run("TerminalRestorer_nil_safe", func(t *testing.T) { - originalRestorer := TerminalRestorer - defer func() { TerminalRestorer = originalRestorer }() - - TerminalRestorer = nil - assert.NotPanics(t, func() { - if TerminalRestorer != nil { - TerminalRestorer() - } - }) - }) -} - -func TestFatalHook(t *testing.T) { - t.Run("fatal_level_calls_TerminalRestorer", func(t *testing.T) { - originalRestorer := TerminalRestorer - defer func() { TerminalRestorer = originalRestorer }() - - called := false - TerminalRestorer = func() { called = true } - - hook := FatalHook{} - - var buf bytes.Buffer - logger := zerolog.New(&buf) - event := logger.Fatal() - - hook.Run(event, zerolog.FatalLevel, "test") - - assert.True(t, called, "TerminalRestorer should be called for fatal level") - }) - - t.Run("non_fatal_level_does_not_call_TerminalRestorer", func(t *testing.T) { - originalRestorer := TerminalRestorer - defer func() { TerminalRestorer = originalRestorer }() - - called := false - TerminalRestorer = func() { called = true } - - hook := FatalHook{} - - levels := []zerolog.Level{ - zerolog.TraceLevel, - zerolog.DebugLevel, - zerolog.InfoLevel, - zerolog.WarnLevel, - zerolog.ErrorLevel, - } - - var buf bytes.Buffer - logger := zerolog.New(&buf) - for _, lvl := range levels { - event := logger.WithLevel(lvl) - hook.Run(event, lvl, "test") - assert.False(t, called, "TerminalRestorer should not be called for non-fatal levels") - } - }) - - t.Run("nil_TerminalRestorer_does_not_panic", func(t *testing.T) { - originalRestorer := TerminalRestorer - defer func() { TerminalRestorer = originalRestorer }() - - TerminalRestorer = nil - hook := FatalHook{} - - assert.NotPanics(t, func() { - var buf bytes.Buffer - logger := zerolog.New(&buf) - event := logger.Fatal() - hook.Run(event, zerolog.FatalLevel, "test") - }) - }) -} - func TestCustomWriter_WritesCorrectly(t *testing.T) { t.Run("Writes_log_to_file", func(t *testing.T) { tmpDir := t.TempDir() diff --git a/pkg/logging/logging.go b/pkg/logging/logging.go index 4f7b4fcf..784d76aa 100644 --- a/pkg/logging/logging.go +++ b/pkg/logging/logging.go @@ -3,8 +3,6 @@ package logging import ( "sync" - "atomicgo.dev/keyboard" - "atomicgo.dev/keyboard/keys" "github.com/rs/zerolog" "github.com/rs/zerolog/log" ) @@ -43,47 +41,27 @@ func defaultStatusHook() *zerolog.Event { } func ShortcutListeners(status ShortcutStatusFN) { - err := keyboard.Listen(func(key keys.Key) (stop bool, err error) { - switch key.Code { - case keys.CtrlC, keys.Escape: - return true, nil - case keys.RuneKey: - if key.String() == "t" { - zerolog.SetGlobalLevel(zerolog.TraceLevel) - log.Info().Str("logLevel", "trace").Msg("New Log level") - } - - if key.String() == "d" { - zerolog.SetGlobalLevel(zerolog.DebugLevel) - log.Info().Str("logLevel", "debug").Msg("New Log level") - } - - if key.String() == "i" { - zerolog.SetGlobalLevel(zerolog.InfoLevel) - log.Info().Str("logLevel", "info").Msg("New Log level") - } - - if key.String() == "w" { - zerolog.SetGlobalLevel(zerolog.WarnLevel) - log.Info().Str("logLevel", "warn").Msg("New Log level") - } - - if key.String() == "e" { - zerolog.SetGlobalLevel(zerolog.ErrorLevel) - log.Info().Str("logLevel", "error").Msg("New Log level") - } - - if key.String() == "s" { - currentHook := GetStatusHook() - log := currentHook() - log.Msg("Status") - } - } - - return false, nil - }) - - if err != nil { + if err := shortcuts.start(); err != nil { log.Error().Err(err).Msg("Failed hooking keyboard bindings") } } + +func handleShortcut(key string) bool { + switch key { + case "\x03", "\x1b": + return true + case "t", "d", "i", "w", "e": + level := map[string]zerolog.Level{ + "t": zerolog.TraceLevel, + "d": zerolog.DebugLevel, + "i": zerolog.InfoLevel, + "w": zerolog.WarnLevel, + "e": zerolog.ErrorLevel, + }[key] + zerolog.SetGlobalLevel(level) + log.Info().Str("logLevel", level.String()).Msg("New Log level") + case "s": + GetStatusHook()().Msg("Status") + } + return false +} diff --git a/pkg/logging/logging_test.go b/pkg/logging/logging_test.go index 6ba2412a..991bf924 100644 --- a/pkg/logging/logging_test.go +++ b/pkg/logging/logging_test.go @@ -3,13 +3,61 @@ package logging import ( "bytes" "io" + "os" "strings" "testing" "github.com/rs/zerolog" "github.com/rs/zerolog/log" + "github.com/stretchr/testify/require" ) +func TestShortcutListenersSkipNonTerminal(t *testing.T) { + reader, writer, err := os.Pipe() + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, reader.Close()) }) + t.Cleanup(func() { require.NoError(t, writer.Close()) }) + originalStdin := os.Stdin + os.Stdin = reader + t.Cleanup(func() { os.Stdin = originalStdin }) + + var terminal shortcutTerminal + require.NoError(t, terminal.start()) + require.Nil(t, terminal.session) + require.NoError(t, terminal.stop()) + require.NoError(t, terminal.stop(), "cleanup must be idempotent") +} + +func TestHandleShortcut(t *testing.T) { + originalLevel := zerolog.GlobalLevel() + originalLogger := log.Logger + originalHook := GetStatusHook() + t.Cleanup(func() { + zerolog.SetGlobalLevel(originalLevel) + log.Logger = originalLogger + RegisterStatusHook(originalHook) + }) + log.Logger = zerolog.New(io.Discard) + for key, level := range map[string]zerolog.Level{ + "t": zerolog.TraceLevel, "d": zerolog.DebugLevel, "i": zerolog.InfoLevel, + "w": zerolog.WarnLevel, "e": zerolog.ErrorLevel, + } { + require.False(t, handleShortcut(key)) + require.Equal(t, level, zerolog.GlobalLevel()) + } + called := false + RegisterStatusHook(func() *zerolog.Event { + called = true + return log.Info() + }) + require.False(t, handleShortcut("s")) + require.True(t, called) + require.False(t, handleShortcut("\x1b[A"), "arrow keys must not stop the listener") + require.False(t, handleShortcut("unknown")) + require.True(t, handleShortcut("\x03")) + require.True(t, handleShortcut("\x1b")) +} + func TestSetLogLevel(t *testing.T) { tests := []struct { name string diff --git a/pkg/logging/terminal.go b/pkg/logging/terminal.go new file mode 100644 index 00000000..8a7d30da --- /dev/null +++ b/pkg/logging/terminal.go @@ -0,0 +1,115 @@ +package logging + +import ( + "errors" + "fmt" + "io" + "os" + "sync" + + "github.com/rs/zerolog" + "github.com/rs/zerolog/log" + "golang.org/x/term" +) + +type shortcutTerminal struct { + mu sync.Mutex + session *shortcutSession +} + +type shortcutSession struct { + input shortcutInput + restore func() error +} + +type shortcutInput interface { + io.ReadWriteCloser + Fd() uintptr + Name() string +} + +var shortcuts shortcutTerminal + +// InstallExitHandler routes fatal logs through the same cleanup as other exits. +func InstallExitHandler() { + zerolog.FatalExitFunc = func() { Exit(1) } +} + +func Exit(code int) { + StopShortcutListeners() + os.Exit(code) +} + +func StopShortcutListeners() { + if err := shortcuts.stop(); err != nil { + log.Error().Err(err).Msg("Failed restoring keyboard terminal") + } +} + +func (t *shortcutTerminal) start() error { + t.mu.Lock() + defer t.mu.Unlock() + if t.session != nil { + return nil + } + // #nosec G115 -- runtime-owned terminal descriptors fit into int. + if !term.IsTerminal(int(os.Stdin.Fd())) { + return nil + } + input, err := openShortcutInput() + if err != nil { + return err + } + restore, err := makeShortcutRaw(input) + if err != nil { + return errors.Join(fmt.Errorf("enable keyboard raw mode: %w", err), input.Close()) + } + session := &shortcutSession{input: input, restore: restore} + t.session = session + go t.listen(session) + return nil +} + +func (t *shortcutTerminal) stop() error { + t.mu.Lock() + defer t.mu.Unlock() + return t.closeLocked() +} + +func (t *shortcutTerminal) closeLocked() error { + if t.session == nil { + return nil + } + session := t.session + t.session = nil + return errors.Join(session.restore(), session.input.Close()) +} + +func (t *shortcutTerminal) listen(session *shortcutSession) { + var buf [256]byte + for { + n, err := session.input.Read(buf[:]) + // Ignore reads unblocked by closing a previous session. + t.mu.Lock() + if t.session != session { + t.mu.Unlock() + return + } + t.mu.Unlock() + if err != nil || handleShortcut(string(buf[:n])) { + t.mu.Lock() + var closeErr error + if t.session == session { + closeErr = t.closeLocked() + } + t.mu.Unlock() + if err != nil && !errors.Is(err, io.EOF) { + log.Error().Err(err).Msg("Failed reading keyboard shortcuts") + } + if closeErr != nil { + log.Error().Err(closeErr).Msg("Failed restoring keyboard terminal") + } + return + } + } +} diff --git a/pkg/logging/terminal_linux_test.go b/pkg/logging/terminal_linux_test.go new file mode 100644 index 00000000..d9a0226a --- /dev/null +++ b/pkg/logging/terminal_linux_test.go @@ -0,0 +1,140 @@ +package logging_test + +import ( + "bytes" + "context" + "errors" + "fmt" + "os" + "os/exec" + "path/filepath" + "strings" + "sync" + "syscall" + "testing" + "time" + + "github.com/CompassSecurity/pipeleek/internal/cmd/common" + "github.com/CompassSecurity/pipeleek/pkg/logging" + "github.com/CompassSecurity/pipeleek/pkg/system" + "github.com/containerd/console" + "github.com/rs/zerolog/log" + "github.com/spf13/cobra" + "github.com/stretchr/testify/require" + "golang.org/x/sys/unix" + "golang.org/x/term" +) + +func TestTerminalLifecycle(t *testing.T) { + for _, mode := range []string{"normal", "error", "fatal-console", "fatal-json", "fatal-logfile", "graceful", "concurrent"} { + t.Run(mode, func(t *testing.T) { + for attempt := 0; attempt < 3; attempt++ { + master, slavePath, err := console.NewPty() + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, master.Close()) }) + slave, err := os.OpenFile(slavePath, os.O_RDWR, 0) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, slave.Close()) }) + // #nosec G115 -- runtime-owned terminal descriptors fit into int. + fd := int(slave.Fd()) + before, err := term.GetState(fd) + require.NoError(t, err) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + child := exec.CommandContext(ctx, os.Args[0], "-test.run=^TestTerminalLifecycleHelper$") + child.Env = append(os.Environ(), "PIPELEEK_TERMINAL_TEST="+mode) + var logPath string + if mode == "fatal-logfile" { + logPath = filepath.Join(t.TempDir(), "fatal.log") + child.Env = append(child.Env, "PIPELEEK_TERMINAL_LOGFILE="+logPath) + } + child.Stdin = slave + child.SysProcAttr = &syscall.SysProcAttr{Setsid: true, Setctty: true, Ctty: 0} + var output bytes.Buffer + child.Stdout = &output + child.Stderr = &output + err = child.Run() + require.NoError(t, ctx.Err(), output.String()) + if mode == "error" || strings.HasPrefix(mode, "fatal-") { + var exitErr *exec.ExitError + require.ErrorAs(t, err, &exitErr, output.String()) + require.Equal(t, 1, exitErr.ExitCode(), output.String()) + } else { + require.NoError(t, err, output.String()) + } + require.Contains(t, output.String(), "RAW_READY") + if mode == "fatal-console" || mode == "fatal-json" { + require.Contains(t, output.String(), "terminal regression fatal") + } + if mode == "fatal-logfile" { + content, err := os.ReadFile(logPath) + require.NoError(t, err) + require.Contains(t, string(content), "terminal regression fatal") + } + after, err := term.GetState(fd) + require.NoError(t, err) + require.Equal(t, before, after, "terminal changed after %s exit", mode) + } + }) + } +} + +func TestTerminalLifecycleHelper(t *testing.T) { + mode := os.Getenv("PIPELEEK_TERMINAL_TEST") + if mode == "" { + return + } + common.JsonLogoutput = mode == "fatal-json" + if mode == "fatal-logfile" { + common.LogFile = os.Getenv("PIPELEEK_TERMINAL_LOGFILE") + } + root := &cobra.Command{Use: "terminal-test", RunE: runTerminalLifecycleHelper} + common.SetupPersistentPreRun(root) + root.SetArgs([]string{}) + common.Run(root) + if mode == "fatal-logfile" { + t.Fatal("fatal log returned instead of exiting") + } +} + +func runTerminalLifecycleHelper(cmd *cobra.Command, args []string) error { + // #nosec G115 -- runtime-owned terminal descriptors fit into int. + state, err := unix.IoctlGetTermios(int(os.Stdin.Fd()), unix.TCGETS) + if err != nil { + return err + } + if state.Lflag&(unix.ECHO|unix.ICANON) != 0 { + return errors.New("shortcut startup returned before raw mode was ready") + } + if state.Oflag&unix.OPOST == 0 { + return errors.New("keyboard raw mode disabled log newline processing") + } + fmt.Println("RAW_READY") + switch mode := os.Getenv("PIPELEEK_TERMINAL_TEST"); { + case mode == "error": + return errors.New("terminal regression command error") + case strings.HasPrefix(mode, "fatal-"): + log.Fatal().Msg("terminal regression fatal") + case mode == "graceful": + system.RegisterGracefulShutdownHandler(func() {}) + if err := syscall.Kill(os.Getpid(), syscall.SIGTERM); err != nil { + return err + } + select {} + case mode == "concurrent": + var wg sync.WaitGroup + for i := 0; i < 20; i++ { + wg.Add(2) + go func() { + defer wg.Done() + logging.ShortcutListeners(nil) + }() + go func() { + defer wg.Done() + logging.StopShortcutListeners() + }() + } + wg.Wait() + } + return nil +} diff --git a/pkg/logging/terminal_unix.go b/pkg/logging/terminal_unix.go new file mode 100644 index 00000000..44b6e08a --- /dev/null +++ b/pkg/logging/terminal_unix.go @@ -0,0 +1,28 @@ +//go:build !windows + +package logging + +import ( + "errors" + "os" + "syscall" + + "github.com/containerd/console" +) + +func openShortcutInput() (shortcutInput, error) { + // Nonblocking open lets Go's poller interrupt the read when cleanup closes it. + return os.OpenFile("/dev/tty", os.O_RDONLY|syscall.O_NONBLOCK, 0) +} + +func makeShortcutRaw(input shortcutInput) (func() error, error) { + c, err := console.ConsoleFromFile(input) + if err != nil { + return nil, err + } + // console.SetRaw preserves output processing, so log newlines stay aligned. + if err := c.SetRaw(); err != nil { + return nil, errors.Join(err, c.Reset()) + } + return c.Reset, nil +} diff --git a/pkg/logging/terminal_windows.go b/pkg/logging/terminal_windows.go new file mode 100644 index 00000000..566c2dee --- /dev/null +++ b/pkg/logging/terminal_windows.go @@ -0,0 +1,86 @@ +package logging + +import ( + "encoding/binary" + "os" + "sync" + "time" + "unsafe" + + "golang.org/x/sys/windows" + "golang.org/x/term" +) + +var readConsoleInput = windows.NewLazySystemDLL("kernel32.dll").NewProc("ReadConsoleInputW") + +type windowsShortcutInput struct { + *os.File + done chan struct{} + mu sync.Mutex +} + +func openShortcutInput() (shortcutInput, error) { + f, err := os.OpenFile("CONIN$", os.O_RDWR, 0) + if err != nil { + return nil, err + } + return &windowsShortcutInput{File: f, done: make(chan struct{})}, nil +} + +func makeShortcutRaw(input shortcutInput) (func() error, error) { + // #nosec G115 -- runtime-owned terminal descriptors fit into int. + fd := int(input.Fd()) + state, err := term.MakeRaw(fd) + if err != nil { + return nil, err + } + return func() error { return term.Restore(fd, state) }, nil +} + +func (input *windowsShortcutInput) Read(p []byte) (int, error) { + input.mu.Lock() + defer input.mu.Unlock() + ticker := time.NewTicker(10 * time.Millisecond) + defer ticker.Stop() + for { + select { + case <-input.done: + return 0, os.ErrClosed + case <-ticker.C: + } + var available uint32 + if err := windows.GetNumberOfConsoleInputEvents(windows.Handle(input.Fd()), &available); err != nil { + return 0, err + } + if available == 0 { + continue + } + // INPUT_RECORD is 20 bytes; its KEY_EVENT_RECORD starts at byte 4. + var record [20]byte + var read uint32 + result, _, err := readConsoleInput.Call( + input.Fd(), uintptr(unsafe.Pointer(&record[0])), 1, uintptr(unsafe.Pointer(&read)), + ) + if result == 0 { + return 0, err + } + if read == 0 || binary.LittleEndian.Uint16(record[:2]) != 1 || + binary.LittleEndian.Uint32(record[4:8]) == 0 { + continue + } + char := binary.LittleEndian.Uint16(record[14:16]) + if char == 0 || char > 127 { + continue + } + return copy(p, []byte{byte(char)}), nil + } +} + +func (input *windowsShortcutInput) Close() error { + // ReadConsole cannot be canceled by os.File.Close. Poll console events instead, + // so cleanup can stop the reader before closing the handle. + close(input.done) + input.mu.Lock() + defer input.mu.Unlock() + return input.File.Close() +} diff --git a/pkg/system/shutdown.go b/pkg/system/shutdown.go index db125384..86f2e832 100644 --- a/pkg/system/shutdown.go +++ b/pkg/system/shutdown.go @@ -5,6 +5,7 @@ import ( "os/signal" "syscall" + "github.com/CompassSecurity/pipeleek/pkg/logging" "github.com/rs/zerolog/log" ) @@ -18,6 +19,6 @@ func RegisterGracefulShutdownHandler(handler ShutdownHandler) { <-sigChannel log.Info().Msg("Received interrupt signal, shutting down gracefully...") handler() - os.Exit(0) + logging.Exit(0) }() }