Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
52 changes: 52 additions & 0 deletions cmd/stepsecurity-dev-machine-guard/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,11 @@ import (
"io"
"os"
"runtime"
"time"

aiagentscli "github.com/step-security/dev-machine-guard/internal/aiagents/cli"
"github.com/step-security/dev-machine-guard/internal/aiagents/ingest"
"github.com/step-security/dev-machine-guard/internal/aiagents/state"
"github.com/step-security/dev-machine-guard/internal/buildinfo"
"github.com/step-security/dev-machine-guard/internal/cli"
"github.com/step-security/dev-machine-guard/internal/config"
Expand All @@ -24,6 +27,12 @@ import (
"github.com/step-security/dev-machine-guard/internal/telemetry"
)

// hookReconcileTimeout caps the entire reconcile step (fetch + cache
// write + install/uninstall). Generous because install can chown a
// handful of files under root; the actual GET cost is bounded by
// state.DefaultFetchTimeout.
const hookReconcileTimeout = 30 * time.Second

func main() {
// Hook hot path. Agents invoke `_hook` on every event and any non-zero
// exit is treated as a hook failure / block — so we MUST exit 0 even on
Expand Down Expand Up @@ -121,6 +130,7 @@ func main() {
log.Error("%v", err)
os.Exit(1)
}
runHookStateReconcile(exec, log)

case "install":
_, _ = fmt.Fprintf(os.Stdout, "StepSecurity Dev Machine Guard v%s\n\n", buildinfo.Version)
Expand Down Expand Up @@ -168,6 +178,7 @@ func main() {
log.Error("%v", telemetryErr)
os.Exit(1)
}
runHookStateReconcile(exec, log)

case "uninstall":
_, _ = fmt.Fprintf(os.Stdout, "StepSecurity Dev Machine Guard v%s\n\n", buildinfo.Version)
Expand Down Expand Up @@ -299,3 +310,44 @@ func scanJSONEncoder(w io.Writer) *json.Encoder {
enc.SetEscapeHTML(false)
return enc
}

// runHookStateReconcile polls agent-api for the desired AI-agent hook
// state and reconciles local hook installation to match. Silent no-op
// in community mode (enterprise config missing) — the existing scan
// path stays unaffected. Failures are logged but never crash main.
func runHookStateReconcile(exec executor.Executor, log *progress.Logger) {
cfg, ok := ingest.Snapshot()
if !ok {
log.Debug("hook-state reconcile: skipped (no enterprise config)")
return
}
fetcher, ok := state.NewHTTPFetcher(cfg, nil)
if !ok {
log.Debug("hook-state reconcile: skipped (fetcher init refused config)")
return
}

ctx, cancel := context.WithTimeout(context.Background(), hookReconcileTimeout)
defer cancel()

dev := device.Gather(ctx, exec)
if dev.SerialNumber == "" || dev.SerialNumber == "unknown" {
log.Warn("hook-state reconcile: device serial unresolved; skipping")
return
}

r := &state.Reconciler{
Exec: exec,
Fetcher: fetcher,
CustomerID: cfg.CustomerID,
DeviceID: dev.SerialNumber,
Stdout: os.Stdout,
Stderr: os.Stderr,
InstallFn: aiagentscli.RunInstall,
UninstallFn: aiagentscli.RunUninstall,
}
if err := r.Reconcile(ctx); err != nil {
log.Warn("hook-state reconcile: %v", err)
aiagentscli.AppendError("reconcile", "reconcile_failed", err.Error(), "")
}
}
11 changes: 11 additions & 0 deletions internal/aiagents/hook/runtime.go
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ import (
"github.com/step-security/dev-machine-guard/internal/aiagents/identity"
"github.com/step-security/dev-machine-guard/internal/aiagents/ingest"
"github.com/step-security/dev-machine-guard/internal/aiagents/redact"
"github.com/step-security/dev-machine-guard/internal/aiagents/state"
"github.com/step-security/dev-machine-guard/internal/executor"
)

Expand Down Expand Up @@ -115,6 +116,16 @@ func (rt *Runtime) Run(parent context.Context, hookType event.HookEvent) error {
var ev *event.Event
defer func() { rt.emitDecidedResponse(ev, decision) }()

// Server-driven kill switch. The reconciler writes ~/.stepsecurity/
// hooks-state.json on every telemetry tick; a UI flip propagates
// to disabled here in O(1) microseconds and short-circuits the
// hot path before any enrichment, identity probe, or upload runs.
// Missing/corrupt cache returns Default()=enabled so first-run
// after install continues to work.
if cur, _ := state.Read(); !cur.Hooks.Enabled {
return nil
}

cfg, _ := ingest.Snapshot()
id := identity.Resolve(ctx, rt.Exec, cfg.CustomerID)
upload := rt.resolveUpload()
Expand Down
44 changes: 44 additions & 0 deletions internal/aiagents/hook/runtime_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,13 +6,15 @@ import (
"encoding/json"
"errors"
"io"
"path/filepath"
"strings"
"sync"
"testing"
"time"

cc "github.com/step-security/dev-machine-guard/internal/aiagents/adapter/claudecode"
"github.com/step-security/dev-machine-guard/internal/aiagents/event"
"github.com/step-security/dev-machine-guard/internal/aiagents/state"
"github.com/step-security/dev-machine-guard/internal/executor"
)

Expand Down Expand Up @@ -420,6 +422,48 @@ func TestRunUploadFailureFailsOpen(t *testing.T) {

// When no UploadEvent is wired (any runtime without enterprise config),
// the runtime must still complete — just with no upload attempt.
// withDisabledStateCache writes a disabled state file and restores
// the override on cleanup. Returns the path for assertions.
func withDisabledStateCache(t *testing.T) {
t.Helper()
dir := t.TempDir()
path := dir + string(filepath.Separator) + state.CacheFilename
restore := state.SetCachePathForTest(path)
t.Cleanup(restore)
s := state.Default()
s.Hooks.Enabled = false
s.Source = state.SourcePoll
if err := state.Write(s); err != nil {
t.Fatalf("seed disabled cache: %v", err)
}
}

func TestRunHonorsDisabledStateCache(t *testing.T) {
withDisabledStateCache(t)

stdin := strings.NewReader(`{
"session_id":"abc","cwd":"/tmp","tool_name":"Bash",
"tool_input":{"command":"npm install lodash","cwd":"/tmp"}
}`)
var stdout, stderr bytes.Buffer
rt, cap := newRuntime(t, stdin, &stdout, &stderr)

if err := rt.Run(context.Background(), event.HookPreToolUse); err != nil {
t.Fatalf("Run: %v", err)
}
// Allow response still emitted (fail-open contract).
if !strings.HasPrefix(strings.TrimSpace(stdout.String()), "{") {
t.Errorf("stdout should still be the allow JSON: %q", stdout.String())
}
// No upload, no enrichment, no error log lines.
if len(cap.events) != 0 {
t.Fatalf("disabled cache must short-circuit upload; got %d events", len(cap.events))
}
if len(cap.errs) != 0 {
t.Fatalf("disabled cache must not log errors; got %v", cap.errs)
}
}

func TestRunSkipsUploadWithoutSeam(t *testing.T) {
stdin := strings.NewReader(`{
"session_id":"s","cwd":"/tmp","tool_name":"Bash","tool_input":{"command":"ls"}
Expand Down
126 changes: 126 additions & 0 deletions internal/aiagents/state/cache.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,126 @@
package state

import (
"encoding/json"
"os"
"path/filepath"
)

// CacheFilename is the basename of the cache file. Lives under
// ~/.stepsecurity/ alongside config.json and ai-agent-hook-errors.jsonl.
const CacheFilename = "hooks-state.json"

const (
cacheFileMode os.FileMode = 0o600
cacheParentDirMode os.FileMode = 0o700
)

// cachePathOverride lets tests redirect reads/writes to a tempdir.
// Production leaves it empty. Mutating from outside this package is
// a test-only concern; same pattern as cli.errorLogPathOverride.
var cachePathOverride string

// SetCachePathForTest redirects CachePath() to the given absolute path
// and returns a restore function. Test-only; production code never
// calls this. Living on the package surface (rather than as a
// build-tagged file) keeps cross-package tests in hook/* and main_test
// able to drive the override without an internal-import trick.
func SetCachePathForTest(p string) (restore func()) {
prev := cachePathOverride
cachePathOverride = p
return func() { cachePathOverride = prev }
}

// CachePath returns the absolute cache path, honoring the test
// override when set.
func CachePath() string {
if cachePathOverride != "" {
return cachePathOverride
}
home, err := os.UserHomeDir()
if err != nil || home == "" {
return ""
}
return filepath.Join(home, ".stepsecurity", CacheFilename)
}

// Read returns (state, true) on a successful parse. Any I/O or parse
// error returns (Default(), false) — Read never surfaces an error
// because the hot path must remain fail-open.
func Read() (State, bool) {
path := CachePath()
if path == "" {
return Default(), false
}
// #nosec G304 -- path is CachePath(): either a test override set by
// SetCachePathForTest, or os.UserHomeDir() joined with the package
// constant CacheFilename. Never derived from external input.
b, err := os.ReadFile(path)
Comment thread
github-advanced-security[bot] marked this conversation as resolved.
Fixed
if err != nil {
return Default(), false
}
var s State
if err := json.Unmarshal(b, &s); err != nil {
return Default(), false
}
if s.SchemaVersion == 0 {
// Forward-compat tolerance: missing schema_version reads as the
// current version. A future breaking change would gate on a
// specific value here.
s.SchemaVersion = SchemaVersion
}
return s, true
}

// Write atomically replaces the cache file. No backups are kept — the
// cache is rewritten on every reconcile tick, and orphaned backups
// would accumulate trash. Parent dir is created with 0o700.
func Write(s State) error {
path := CachePath()
if path == "" {
return errNoHomeDir
}
data, err := json.MarshalIndent(s, "", " ")
if err != nil {
return err
}
data = append(data, '\n')

parent := filepath.Dir(path)
if err := os.MkdirAll(parent, cacheParentDirMode); err != nil {
return err
}

tmp, err := os.CreateTemp(parent, "."+CacheFilename+".tmp-*")
if err != nil {
return err
}
tmpPath := tmp.Name()
defer func() {
if _, statErr := os.Stat(tmpPath); statErr == nil {
_ = os.Remove(tmpPath)
}
}()

if _, err := tmp.Write(data); err != nil {
_ = tmp.Close()
return err
}
if err := tmp.Sync(); err != nil {
_ = tmp.Close()
return err
}
if err := tmp.Close(); err != nil {
return err
}
if err := os.Chmod(tmpPath, cacheFileMode); err != nil {
return err
}
return os.Rename(tmpPath, path)
}

type cacheError string

func (e cacheError) Error() string { return string(e) }

const errNoHomeDir = cacheError("state: cannot resolve home directory")
Loading
Loading