diff --git a/README.md b/README.md index a999b3d8..8bb636ff 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,7 @@

aiscan logo

aiscan

-

Agentic Security Scanner — AI-driven reconnaissance meets deterministic scanning

+

AI-driven single-binary pentest agent with a built-in multi-engine arsenal, ready to go

Preview — APIs and features may change between releases

diff --git a/README_CN.md b/README_CN.md index f8b46453..baea6796 100644 --- a/README_CN.md +++ b/README_CN.md @@ -1,7 +1,7 @@

aiscan logo

aiscan

-

Agentic Security Scanner — AI 驱动的侦察与确定性扫描融合

+

AI 驱动的面向实战的单文件渗透 agent，内置多引擎武器库开箱即用

Preview — 本项目处于早期预览阶段，API 和功能可能随版本变更

diff --git a/build.sh b/build.sh index 44ae8b70..ac6e5b66 100755 --- a/build.sh +++ b/build.sh @@ -177,7 +177,7 @@ CFG_IOA_SPACE=$(resolve "$OPT_IOA_SPACE" "$(yaml_val "$CONFIG_FILE" ioa space)") CFG_VERIFY=$(resolve "$OPT_VERIFY" "$(yaml_val "$CONFIG_FILE" scan verify)") CFG_VERIFY_TIMEOUT=$(resolve "$OPT_VERIFY_TIMEOUT" "$(yaml_val "$CONFIG_FILE" scan verify_timeout)") -CFG_TAVILY_KEYS=$(resolve "$OPT_TAVILY_KEYS" "$(yaml_val "$CONFIG_FILE" websearch tavily_keys)") +CFG_TAVILY_KEYS=$(resolve "$OPT_TAVILY_KEYS" "$(yaml_val "$CONFIG_FILE" search tavily_keys)") # build 段仅从 aiscan.yaml 读取（不做 CLI 覆盖） if [ -z "$OSARCH" ]; then diff --git a/cmd/aiscan/imports.go b/cmd/aiscan/imports.go index aed85fea..ac30b23f 100644 --- a/cmd/aiscan/imports.go +++ b/cmd/aiscan/imports.go @@ -6,8 +6,12 @@ package main import ( _ "github.com/chainreactors/aiscan/pkg/tools" _ "github.com/chainreactors/aiscan/pkg/tools/arsenal" - _ "github.com/chainreactors/aiscan/pkg/tools/proton" + _ "github.com/chainreactors/aiscan/pkg/tools/gogo" _ "github.com/chainreactors/aiscan/pkg/tools/ioa" + _ "github.com/chainreactors/aiscan/pkg/tools/neutron" + _ "github.com/chainreactors/aiscan/pkg/tools/proton" _ "github.com/chainreactors/aiscan/pkg/tools/proxy" _ "github.com/chainreactors/aiscan/pkg/tools/search" + _ "github.com/chainreactors/aiscan/pkg/tools/spray" + _ "github.com/chainreactors/aiscan/pkg/tools/zombie" ) diff --git a/core/config/config_test.go b/core/config/config_test.go index 14485425..59459e35 100644 --- a/core/config/config_test.go +++ b/core/config/config_test.go @@ -344,7 +344,7 @@ llm: } if path == "" { - t.Fatal("expected config.yaml to be found") + t.Fatal("expected aiscan.yaml to be found") } if option.Provider != "found-provider" { t.Errorf("Provider: got %q, want %q", option.Provider, "found-provider") diff --git a/docs/changelog.md b/docs/changelog.md index 944caaae..aaa62d25 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -1,8 +1,8 @@ # Changelog -## v0.2.7 — Proton 敏感信息扫描 + /loop 循环任务 + TUI 交互增强 + 多 Provider 配置 +## v0.2.7 — MITM 流量捕获 + Proton 敏感信息扫描 + /loop 循环任务 + TUI 交互增强 -新增 Proton 敏感信息扫描器（SDK 引擎 + 197 条内嵌规则 + 双向管道）；`/loop` 循环任务调度；TUI 交互全面增强（verbosity 切换、中断控制、文件补全、实时 token 用量）；并发工具执行 OOM 防护；多 Provider 列表配置格式；FOFA key-only 认证支持。 +MITM 透明流量拦截（`proxy mitm` 子命令族）；Proton 敏感信息扫描器（SDK 引擎 + 197 条内嵌规则 + 双向管道）；`/loop` 循环任务调度；TUI 交互全面增强（verbosity 切换、中断控制、文件补全、实时 token 用量）；多 Provider 列表配置格式；FOFA key-only 认证支持。 ### New Features @@ -67,6 +67,36 @@ bash(command="loop */5 * * * * check scan progress") bash(command="loop list") ``` +**MITM 流量捕获** + +透明 HTTP/HTTPS 流量拦截，集成 utils/mitmproxy 到 proxy 命令组。扫描引擎（gogo/spray/zombie/neutron）自动路由到本地 MITM 代理；若已有外部代理（trojan/vless/clash）则作为上游透传。 + +- `proxy mitm start [--addr]`：启动本地 MITM 代理，自动切换扫描引擎代理 +- `proxy mitm stop`：停止 MITM 并恢复之前的代理设置 +- `proxy mitm status`：查看状态和 flow 计数 +- `proxy mitm flows [--host/--status/--type/--last]`：按条件查询捕获的 HTTP 流 +- `proxy mitm flow `：查看单个 flow 详情 +- `proxy mitm clear`：清空 flow 存储 +- `proxy mitm analyze [--host/--last]`：结构化输出供 AI 分析 + +```bash +# 启动 MITM 拦截 +proxy mitm start --addr 127.0.0.1:8888 + +# 正常执行扫描（流量自动经过 MITM） +scan -i target + +# 查看捕获的流量 +proxy mitm flows --last 20 +proxy mitm flow 42 + +# AI 分析捕获的请求 +proxy mitm analyze --host target.com + +# 停止并恢复 +proxy mitm stop +``` + ### Improvements **TUI 交互增强** @@ -112,8 +142,6 @@ llm: - **Resources 统一**：4 套独立 config map（gogo/spray/zombie/proton）合并为单一 `configs map[string]map[string][]byte` + `Config(engine, name)` 方法 - **toolargs 共享工具包**：提取 `ResolveRelativePaths` 和 `NormalizeFlags` 到 `toolargs/`，6 个工具共用（proton/neutron/scan/spray/zombie/katana） -- **Deps.GetLogger()**：消除 5 个 register.go 中重复的 logger 提取模式 -- 净减 180 行代码 ### Bug Fixes @@ -127,6 +155,7 @@ llm: ### Dependencies - **spray [v1.3.1](https://github.com/chainreactors/spray/releases/tag/v1.3.1)**：mask 表达式支持所有请求字段、`--keys` 插件内嵌 156 条 proton 模板、extract severity 分级 + 上下文捕获、修复 crawl-only 提前 drain 和 OutputCh panic +- **utils/cert**：集成 utils/cert 原子化证书原语（CA 生成、子证书签发、随机 Subject、PEM 工具函数），移除本地 replace - bump SDK、zombie、logs、utils/pty 修复上游 data race ### Breaking Changes diff --git a/go.mod b/go.mod index aa0be0fe..a2320f5a 100644 --- a/go.mod +++ b/go.mod @@ -18,7 +18,7 @@ require ( github.com/chainreactors/sdk v0.3.4-0.20260624031614-b16da9a87441 github.com/chainreactors/spray v1.3.2-0.20260624034433-890630649b2b github.com/chainreactors/utils v0.0.0-20260623065725-737b33d61c6b - github.com/chainreactors/utils/mitmproxy v0.0.0 + github.com/chainreactors/utils/mitmproxy v0.0.0-20260624182357-8d5cad72d8f2 github.com/chainreactors/utils/pty v0.0.0-20260624031611-9aadeae3fb0e github.com/chainreactors/zombie v1.2.3-0.20260624041317-6bf4579de29d github.com/charmbracelet/bubbles v1.0.0 @@ -88,7 +88,7 @@ require ( github.com/censys/censys-sdk-go v0.19.1 // indirect github.com/chainreactors/files v0.0.0-20240716182835-7884ee1e77f0 // indirect github.com/chainreactors/neutron/operators/full v0.0.0-20260615055126-a9bbe4fc3e95 // indirect - github.com/chainreactors/utils/cert v0.0.0 // indirect + github.com/chainreactors/utils/cert v0.0.0-20260624181253-2b3d0b35862f // indirect github.com/chainreactors/words v0.0.0-20260520145736-270600e60fb4 // indirect github.com/charlievieth/fastwalk v1.0.14 // indirect github.com/charmbracelet/bubbletea v1.3.10 // indirect @@ -308,8 +308,6 @@ require ( ) replace ( - github.com/chainreactors/utils/cert => ../utils/cert - github.com/chainreactors/utils/mitmproxy => ../utils/mitmproxy github.com/reeflective/console => github.com/chainreactors/malice-network/external/console v0.0.0-20260620124902-e62d93c864f7 github.com/reeflective/readline => github.com/chainreactors/malice-network/external/readline v0.0.0-20260620124902-e62d93c864f7 ) diff --git a/go.sum b/go.sum index 81afaa99..1cabf826 100644 --- a/go.sum +++ b/go.sum @@ -206,6 +206,10 @@ github.com/chainreactors/spray v1.3.2-0.20260624034433-890630649b2b/go.mod h1:uh github.com/chainreactors/utils v0.0.0-20240716182459-e85f2b01ee16/go.mod h1:LajXuvESQwP+qCMAvlcoSXppQCjuLlBrnQpu9XQ1HtU= github.com/chainreactors/utils v0.0.0-20260623065725-737b33d61c6b h1:JWciIo6GXn+QqejE8ER0oo/IxqaAAlSaIqsHhZtnoMM= github.com/chainreactors/utils v0.0.0-20260623065725-737b33d61c6b/go.mod h1:LajXuvESQwP+qCMAvlcoSXppQCjuLlBrnQpu9XQ1HtU= +github.com/chainreactors/utils/cert v0.0.0-20260624181253-2b3d0b35862f h1:7KonUAvkZhpmK8RDwDBCuey0qXeoqxJ7uWjoXqzL2z4= +github.com/chainreactors/utils/cert v0.0.0-20260624181253-2b3d0b35862f/go.mod h1:xvvWMcU9Fcht6GR1cc9ceAZ3/Hl2HrkoRzpeyOzx1rQ= +github.com/chainreactors/utils/mitmproxy v0.0.0-20260624182357-8d5cad72d8f2 h1:9sZzzfeYX5gD64TEpCykenqY+fhej6g2s/lCzMDNNIg= +github.com/chainreactors/utils/mitmproxy v0.0.0-20260624182357-8d5cad72d8f2/go.mod h1:2O3/Vw66VnbzhwsHGFJ2Ge98RuSh6XzMMFGZmMmlZ9M= github.com/chainreactors/utils/pty v0.0.0-20260624031611-9aadeae3fb0e h1:KzDf08J9nGhTjyTz1pMR7aMnkbTEeagZllFcZnZuthM= github.com/chainreactors/utils/pty v0.0.0-20260624031611-9aadeae3fb0e/go.mod h1:RW1v+8hFMeO9+TJyQ1iIx9Ea37s+B7BaDf9fyJ0OEC4= github.com/chainreactors/words v0.0.0-20260520145736-270600e60fb4 h1:lvnDYEkatmZFHP5i321qQXK9L4vKRfso/uUfr5tOeC8= diff --git a/pkg/agent/context_window.go b/pkg/agent/context_window.go index c6de3817..c2fac926 100644 --- a/pkg/agent/context_window.go +++ b/pkg/agent/context_window.go @@ -34,6 +34,8 @@ var modelContextWindows = []struct { {"gpt-4.1", 1047576}, {"gpt-4o", 128000}, {"gpt-4-turbo", 128000}, + {"gpt-4-1", 128000}, + {"gpt-4-0", 8192}, {"gpt-4", 8192}, {"o4-mini", 200000}, {"o3", 200000}, diff --git a/pkg/agent/evaluator/evaluator.go b/pkg/agent/evaluator/evaluator.go index d12cdc78..b7af74ba 100644 --- a/pkg/agent/evaluator/evaluator.go +++ b/pkg/agent/evaluator/evaluator.go @@ -65,7 +65,11 @@ func (e *Evaluator) Evaluate(ctx context.Context, goal, criteria string, message lastErr = err e.cfg.Logger.Warnf("evaluate attempt %d failed: %s", attempt+1, err) if attempt < e.cfg.MaxRetries-1 { - time.Sleep(time.Duration(attempt+1) * time.Second) + select { + case <-time.After(time.Duration(attempt+1) * time.Second): + case <-ctx.Done(): + return nil, ctx.Err() + } } } return nil, fmt.Errorf("evaluate failed after %d attempts: %w", e.cfg.MaxRetries, lastErr) diff --git a/pkg/agent/loop_scheduler.go b/pkg/agent/loop_scheduler.go index f2def9b3..baff77a7 100644 --- a/pkg/agent/loop_scheduler.go +++ b/pkg/agent/loop_scheduler.go @@ -70,6 +70,7 @@ type LoopScheduler struct { type loopState struct { entry LoopEntry cancel context.CancelFunc + wg sync.WaitGroup fireCount int lastFired time.Time } @@ -198,7 +199,9 @@ func (s *LoopScheduler) fire(ctx context.Context, state *loopState) { } case ModeIndependent: + state.wg.Add(1) go func() { + defer state.wg.Done() if _, err := entry.OnFire(ctx, entry); err != nil { s.log.Warnf("loop=%s fire=%d failed: %s", entry.Name, count, err) } @@ -216,6 +219,7 @@ func (s *LoopScheduler) Remove(name string) error { state.cancel() delete(s.loops, name) s.mu.Unlock() + state.wg.Wait() s.log.Importantf("loop=%s deleted", name) return nil } @@ -245,9 +249,14 @@ func (s *LoopScheduler) Active() int { func (s *LoopScheduler) Stop() { s.mu.Lock() - defer s.mu.Unlock() + states := make([]*loopState, 0, len(s.loops)) for name, state := range s.loops { state.cancel() + states = append(states, state) delete(s.loops, name) } + s.mu.Unlock() + for _, state := range states { + state.wg.Wait() + } } diff --git a/pkg/agent/loop_tool.go b/pkg/agent/loop_tool.go index 0afa237a..8751dc9a 100644 --- a/pkg/agent/loop_tool.go +++ b/pkg/agent/loop_tool.go @@ -102,6 +102,11 @@ func (c *LoopCommand) create(ctx context.Context, args []string) error { // tryCronPrefix attempts to parse the first 5 args as a cron expression. // Returns the parsed expression, remaining args, and whether it succeeded. func tryCronPrefix(args []string) (*CronExpr, []string, bool) { + if len(args) >= 2 && strings.Contains(args[0], " ") { + if cron, err := ParseCron(args[0]); err == nil { + return cron, args[1:], true + } + } if len(args) < 6 { return nil, nil, false } diff --git a/pkg/agent/types.go b/pkg/agent/types.go index 7f3a180e..4bec4ee9 100644 --- a/pkg/agent/types.go +++ b/pkg/agent/types.go @@ -232,7 +232,7 @@ func (c Config) init() Config { if c.Logger == nil { c.Logger = telemetry.NopLogger() } - if c.MaxRetries == 0 { + if c.MaxRetries < 0 { c.MaxRetries = DefaultMaxRetries } if c.MaxResultSize <= 0 { diff --git a/pkg/commands/factory.go b/pkg/commands/factory.go index 5c8ddcaa..3b720c2f 100644 --- a/pkg/commands/factory.go +++ b/pkg/commands/factory.go @@ -22,17 +22,15 @@ type Deps struct { Provider any ScannerProxy string ScanOpts []any - Logger any + Logger telemetry.Logger NodeName string NodeMeta map[string]any TavilyKeys string // comma-separated Tavily API keys (build-time fallback) } func (d *Deps) GetLogger() telemetry.Logger { - if d != nil { - if logger, ok := d.Logger.(telemetry.Logger); ok && logger != nil { - return logger - } + if d != nil && d.Logger != nil { + return d.Logger } return telemetry.NopLogger() } diff --git a/pkg/tools/gogo/gogo.go b/pkg/tools/gogo/gogo.go index dd41ccf9..35f7df7e 100644 --- a/pkg/tools/gogo/gogo.go +++ b/pkg/tools/gogo/gogo.go @@ -15,32 +15,26 @@ import ( ) type Command struct { - engine *gogo.Engine - logger telemetry.Logger - proxy string - workDir string + toolargs.Base + engine *gogo.Engine } func New(engine *gogo.Engine) *Command { - return &Command{engine: engine, logger: telemetry.NopLogger()} + c := &Command{engine: engine} + c.InitLogger(nil) + return c } func (c *Command) WithLogger(logger telemetry.Logger) *Command { - if logger != nil { - c.logger = logger - } + c.InitLogger(logger) return c } -func (c *Command) SetWorkDir(dir string) { c.workDir = dir } - func (c *Command) WithProxy(proxy string) *Command { - c.proxy = proxy + c.Proxy = proxy return c } -func (c *Command) SetProxy(proxy string) { c.proxy = proxy } - func (c *Command) Name() string { return "gogo" } func (c *Command) Usage() string { @@ -52,9 +46,9 @@ func (c *Command) Execute(ctx context.Context, args []string) (err error) { args = c.injectProxy(args) if toolargs.BoolFlagEnabled(args, "--debug") { - restoreDebug := telemetry.ActivateDebug(c.logger) + restoreDebug := telemetry.ActivateDebug(c.Logger) defer restoreDebug() - c.logger.Debugf("gogo debug enabled") + c.Logger.Debugf("gogo debug enabled") } var buf bytes.Buffer @@ -86,13 +80,13 @@ func (c *Command) TestInjectProxy(args []string) []string { } func (c *Command) injectProxy(args []string) []string { - if c.proxy == "" { + if c.Proxy == "" { return args } if toolargs.HasFlag(args, "--proxy") { return args } - return append(args, "--proxy", c.proxy) + return append(args, "--proxy", c.Proxy) } // normalizeArgs adapts common agent-generated gogo arguments before handing @@ -139,10 +133,10 @@ func (c *Command) normalizeArgs(args []string) []string { } func (c *Command) resolvePathArg(value string) string { - if c.workDir == "" || value == "" || filepath.IsAbs(value) || strings.HasPrefix(value, "-") { + if c.WorkDir == "" || value == "" || filepath.IsAbs(value) || strings.HasPrefix(value, "-") { return value } - return filepath.Join(c.workDir, value) + return filepath.Join(c.WorkDir, value) } func isGogoValuelessJSONFlag(arg string, args []string, index int) bool { diff --git a/pkg/tools/gogo/register.go b/pkg/tools/gogo/register.go new file mode 100644 index 00000000..921c29eb --- /dev/null +++ b/pkg/tools/gogo/register.go @@ -0,0 +1,22 @@ +package gogo + +import ( + "github.com/chainreactors/aiscan/pkg/commands" + "github.com/chainreactors/aiscan/pkg/tools/scan/engine" +) + +func init() { + commands.RegisterFactory(commands.Factory{ + Group: "scanner", + Build: func(deps *commands.Deps, reg *commands.CommandRegistry) { + es, _ := deps.EngineSet.(*engine.Set) + if es == nil || es.Gogo == nil { + return + } + reg.Register( + New(es.Gogo).WithLogger(deps.GetLogger()).WithProxy(deps.ScannerProxy), + "scanner", + ) + }, + }) +} diff --git a/pkg/tools/katana/katana.go b/pkg/tools/katana/katana.go index 63f3f0ca..f9d8f07a 100644 --- a/pkg/tools/katana/katana.go +++ b/pkg/tools/katana/katana.go @@ -30,29 +30,25 @@ const defaultTimeout = 120 * time.Second const defaultBodyReadSize = 4 * 1024 * 1024 type Command struct { - proxy string - logger telemetry.Logger - workDir string + toolargs.Base } -func New() *Command { return &Command{logger: telemetry.NopLogger()} } +func New() *Command { + c := &Command{} + c.InitLogger(nil) + return c +} func (c *Command) WithLogger(logger telemetry.Logger) *Command { - if logger != nil { - c.logger = logger - } + c.InitLogger(logger) return c } func (c *Command) WithProxy(proxy string) *Command { - c.proxy = proxy + c.Proxy = proxy return c } -func (c *Command) SetProxy(proxy string) { c.proxy = proxy } - -func (c *Command) SetWorkDir(dir string) { c.workDir = dir } - func (c *Command) Name() string { return "katana" } func (c *Command) Usage() string { @@ -115,9 +111,9 @@ func (c *Command) Execute(ctx context.Context, args []string) (err error) { args = c.resolveRelativePaths(args) if toolargs.BoolFlagEnabled(args, "--debug") { - restoreDebug := telemetry.ActivateDebug(c.logger) + restoreDebug := telemetry.ActivateDebug(c.Logger) defer restoreDebug() - c.logger.Debugf("katana debug enabled") + c.Logger.Debugf("katana debug enabled") } options, err := readFlags(args) @@ -131,8 +127,8 @@ func (c *Command) Execute(ctx context.Context, args []string) (err error) { options.DisableUpdateCheck = true // Inject proxy. - if options.Proxy == "" && c.proxy != "" { - options.Proxy = c.proxy + if options.Proxy == "" && c.Proxy != "" { + options.Proxy = c.Proxy } if err := validateOptions(options); err != nil { @@ -197,7 +193,7 @@ func (c *Command) Execute(ctx context.Context, args []string) (err error) { if ctx.Err() != nil { return fmt.Errorf("katana: timed out") } - c.logger.Warnf("katana: crawl %s: %v", u, crawlErr) + c.Logger.Warnf("katana: crawl %s: %v", u, crawlErr) } } @@ -368,7 +364,7 @@ func validateOptions(options *katanatypes.Options) error { } func (c *Command) resolveRelativePaths(args []string) []string { - if c.workDir == "" { + if c.WorkDir == "" { return args } out := make([]string, 0, len(args)) @@ -392,10 +388,10 @@ func (c *Command) resolveRelativePaths(args []string) []string { } func (c *Command) resolvePathArg(value string) string { - if c.workDir == "" || value == "" || filepath.IsAbs(value) || strings.HasPrefix(value, "-") { + if c.WorkDir == "" || value == "" || filepath.IsAbs(value) || strings.HasPrefix(value, "-") { return value } - return filepath.Join(c.workDir, value) + return filepath.Join(c.WorkDir, value) } func isFileFlag(flag string) bool { diff --git a/pkg/tools/neutron/neutron.go b/pkg/tools/neutron/neutron.go index 7b5427b9..6e0802c2 100644 --- a/pkg/tools/neutron/neutron.go +++ b/pkg/tools/neutron/neutron.go @@ -23,15 +23,11 @@ import ( ) type Command struct { - engine *sdkneutron.Engine - index *association.Index - logger telemetry.Logger - proxy string - workDir string + toolargs.Base + engine *sdkneutron.Engine + index *association.Index } -func (c *Command) SetWorkDir(dir string) { c.workDir = dir } - type neutronFlags struct { Inputs []string `short:"u" long:"target" description:"Target URL, host, or ip:port (can specify multiple)"` Input string `short:"i" long:"input" description:"Target URL, host, or ip:port (alias of --target)"` @@ -84,24 +80,24 @@ type neutronSummary struct { } func New(engine *sdkneutron.Engine, index *association.Index) *Command { - return &Command{engine: engine, index: index, logger: telemetry.NopLogger()} + c := &Command{engine: engine, index: index} + c.InitLogger(nil) + return c } func (c *Command) WithLogger(logger telemetry.Logger) *Command { - if logger != nil { - c.logger = logger - } + c.InitLogger(logger) return c } func (c *Command) WithProxy(proxy string) *Command { - c.proxy = proxy + c.Proxy = proxy scanengine.ApplyNeutronProxy(proxy) return c } func (c *Command) SetProxy(proxy string) { - c.proxy = proxy + c.Base.SetProxy(proxy) scanengine.ApplyNeutronProxy(proxy) } @@ -159,9 +155,9 @@ func (c *Command) Execute(ctx context.Context, args []string) error { return fmt.Errorf("neutron: %w", err) } if flags.Debug { - restoreDebug := telemetry.ActivateDebug(c.logger) + restoreDebug := telemetry.ActivateDebug(c.Logger) defer restoreDebug() - c.logger.Debugf("neutron debug enabled") + c.Logger.Debugf("neutron debug enabled") } if flags.Timeout > 0 { var cancel context.CancelFunc @@ -232,7 +228,7 @@ func (c *Command) Execute(ctx context.Context, args []string) error { return wErr } - c.logger.Infof("neutron action=testing targets=%d templates=%d concurrency=%d rate_limit=%d", len(targets), len(selected), flags.Concurrency, flags.RateLimit) + c.Logger.Infof("neutron action=testing targets=%d templates=%d concurrency=%d rate_limit=%d", len(targets), len(selected), flags.Concurrency, flags.RateLimit) summary := neutronSummary{Targets: len(targets), Templates: len(selected)} var sb strings.Builder jsonOutput := flags.JSON || flags.JSONL @@ -526,7 +522,7 @@ var neutronFileFlags = map[string]bool{ } func (c *Command) resolveRelativePaths(args []string) []string { - return toolargs.ResolveRelativePaths(args, neutronFileFlags, c.workDir) + return toolargs.ResolveRelativePaths(args, neutronFileFlags, c.WorkDir) } func appendNonEmpty(parts []string, values ...string) []string { diff --git a/pkg/tools/neutron/register.go b/pkg/tools/neutron/register.go new file mode 100644 index 00000000..4a24ef35 --- /dev/null +++ b/pkg/tools/neutron/register.go @@ -0,0 +1,22 @@ +package neutron + +import ( + "github.com/chainreactors/aiscan/pkg/commands" + "github.com/chainreactors/aiscan/pkg/tools/scan/engine" +) + +func init() { + commands.RegisterFactory(commands.Factory{ + Group: "scanner", + Build: func(deps *commands.Deps, reg *commands.CommandRegistry) { + es, _ := deps.EngineSet.(*engine.Set) + if es == nil || es.Neutron == nil { + return + } + reg.Register( + New(es.Neutron, es.Index).WithLogger(deps.GetLogger()).WithProxy(deps.ScannerProxy), + "scanner", + ) + }, + }) +} diff --git a/pkg/tools/proton/command.go b/pkg/tools/proton/command.go index e71a25f7..b6062d78 100644 --- a/pkg/tools/proton/command.go +++ b/pkg/tools/proton/command.go @@ -25,26 +25,24 @@ import ( ) type Command struct { - logger telemetry.Logger - proxy string - workDir string + toolargs.Base stdinFile string resourceProvider func(string) []byte } func New() *Command { - return &Command{logger: telemetry.NopLogger()} + c := &Command{} + c.InitLogger(nil) + return c } func (c *Command) WithLogger(logger telemetry.Logger) *Command { - if logger != nil { - c.logger = logger - } + c.InitLogger(logger) return c } func (c *Command) WithProxy(proxy string) *Command { - c.proxy = proxy + c.Proxy = proxy return c } @@ -53,8 +51,6 @@ func (c *Command) WithResourceProvider(provider func(string) []byte) *Command { return c } -func (c *Command) SetProxy(proxy string) { c.proxy = proxy } -func (c *Command) SetWorkDir(dir string) { c.workDir = dir } func (c *Command) SetStdinFile(path string) { c.stdinFile = path } func (c *Command) Name() string { return "proton" } @@ -152,9 +148,9 @@ func (c *Command) Execute(ctx context.Context, args []string) error { } if flags.Debug { - restoreDebug := telemetry.ActivateDebug(c.logger) + restoreDebug := telemetry.ActivateDebug(c.Logger) defer restoreDebug() - c.logger.Debugf("proton debug enabled") + c.Logger.Debugf("proton debug enabled") } if flags.Timeout > 0 { var cancel context.CancelFunc @@ -213,10 +209,13 @@ func (c *Command) Execute(ctx context.Context, args []string) error { if len(inputs) == 0 && len(flags.Expressions) == 0 { return fmt.Errorf("proton: target required (-i , -l , -e , or pipe: | proton)") } + if len(inputs) == 0 { + inputs = []string{"."} + } // --- Scan --- sevFilter := buildSeverityFilter(flags.Severity, flags.ExcludeSeverity) - seen := make(map[string]bool) + var seen sync.Map var findingCount int64 var fileOut *os.File @@ -234,10 +233,9 @@ func (c *Command) Execute(ctx context.Context, args []string) error { return } key := uf.TemplateID + "|" + uf.FilePath - if seen[key] { + if _, loaded := seen.LoadOrStore(key, true); loaded { return } - seen[key] = true atomic.AddInt64(&findingCount, 1) writeFinding(commands.Output, uf, flags.JSON, inputs[0]) if fileOut != nil { @@ -245,12 +243,12 @@ func (c *Command) Execute(ctx context.Context, args []string) error { } } - c.logger.Infof("proton action=scanning targets=%d rules=%d", len(inputs), scanner.Stats.Rules) + c.Logger.Infof("proton action=scanning targets=%d rules=%d", len(inputs), scanner.Stats.Rules) for _, input := range inputs { info, statErr := os.Stat(input) if statErr != nil { - c.logger.Warnf("proton: skip %s: %v", input, statErr) + c.Logger.Warnf("proton: skip %s: %v", input, statErr) continue } if info.IsDir() { @@ -353,7 +351,7 @@ var protonFileFlags = map[string]bool{ } func (c *Command) resolveRelativePaths(args []string) []string { - return toolargs.ResolveRelativePaths(args, protonFileFlags, c.workDir) + return toolargs.ResolveRelativePaths(args, protonFileFlags, c.WorkDir) } // --- output --- @@ -382,8 +380,9 @@ func writeFinding(w interface{ Write([]byte) (int, error) }, f file.Finding, jso } func truncate(s string, max int) string { - if len(s) > max { - return s[:max] + "..." + runes := []rune(s) + if len(runes) > max { + return string(runes[:max]) + "..." } return s } @@ -444,8 +443,11 @@ func walkAndScan(ctx context.Context, scanner *file.Scanner, target string, call }() } - _ = filepath.WalkDir(target, func(path string, d fs.DirEntry, err error) error { - if err != nil || ctx.Err() != nil { + if walkErr := filepath.WalkDir(target, func(path string, d fs.DirEntry, err error) error { + if ctx.Err() != nil { + return ctx.Err() + } + if err != nil { return err } if d.IsDir() { @@ -468,7 +470,9 @@ func walkAndScan(ctx context.Context, scanner *file.Scanner, target string, call jobCh <- job{path: path, group: group} } return nil - }) + }); walkErr != nil && ctx.Err() == nil { + fmt.Fprintf(os.Stderr, "proton: walk %s: %v\n", target, walkErr) + } close(jobCh) wg.Wait() } diff --git a/pkg/tools/proxy/mitm.go b/pkg/tools/proxy/mitm.go index 14fb7540..3c521031 100644 --- a/pkg/tools/proxy/mitm.go +++ b/pkg/tools/proxy/mitm.go @@ -4,14 +4,12 @@ import ( "context" "fmt" "net/http" - "net/url" "strconv" "strings" "sync" "time" "github.com/chainreactors/aiscan/pkg/commands" - "github.com/chainreactors/proxyclient" mitmproxy "github.com/chainreactors/utils/mitmproxy/proxy" goflags "github.com/jessevdk/go-flags" ) @@ -182,7 +180,6 @@ func (s *mitmState) start() error { if err != nil { return fmt.Errorf("create MITM proxy: %w", err) } - s.store.Clear() p.AddAddon(&captureAddon{store: s.store}) listenAddr, _, err := p.StartAsync() if err != nil { @@ -196,7 +193,7 @@ func (s *mitmState) start() error { func (s *mitmState) stop() { if s.server != nil { ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) - s.server.Shutdown(ctx) + _ = s.server.Shutdown(ctx) cancel() s.server = nil } @@ -225,7 +222,9 @@ func (a *captureAddon) Requestheaders(f *mitmproxy.Flow) { func (a *captureAddon) Response(f *mitmproxy.Flow) { var dur time.Duration if start, ok := a.pending.LoadAndDelete(f.Id.String()); ok { - dur = time.Since(start.(time.Time)) + if t, ok := start.(time.Time); ok { + dur = time.Since(t) + } } flow := Flow{ Timestamp: f.StartTime, @@ -253,7 +252,9 @@ func (a *captureAddon) Response(f *mitmproxy.Flow) { func (a *captureAddon) RequestError(f *mitmproxy.Flow, err error) { var dur time.Duration if start, ok := a.pending.LoadAndDelete(f.Id.String()); ok { - dur = time.Since(start.(time.Time)) + if t, ok := start.(time.Time); ok { + dur = time.Since(t) + } } a.store.Add(Flow{ Timestamp: f.StartTime, @@ -368,6 +369,7 @@ func (s *FlowStore) Clear() { s.mu.Lock() defer s.mu.Unlock() s.flows = s.flows[:0] + s.seq = 0 } func (s *FlowStore) Count() int { @@ -393,7 +395,7 @@ func matchStatus(code int, pattern string) bool { if n, err := strconv.Atoi(p); err == nil { return code == n } - return true + return false } } @@ -499,20 +501,3 @@ func writeHeaders(sb *strings.Builder, h http.Header) { } } } - -// configureUpstream sets the MITM proxy's upstream dialer if a proxy is active. -func configureUpstream(p *mitmproxy.Proxy, proxyURL string) error { - if proxyURL == "" { - return nil - } - u, err := url.Parse(proxyURL) - if err != nil { - return err - } - dial, err := proxyclient.NewClient(u) - if err != nil { - return fmt.Errorf("upstream dialer for %s: %w", u.Scheme, err) - } - p.SetDialer(dial.DialContext) - return nil -} diff --git a/pkg/tools/proxy/mitm_bench_test.go b/pkg/tools/proxy/mitm_bench_test.go index 69d19444..491cf125 100644 --- a/pkg/tools/proxy/mitm_bench_test.go +++ b/pkg/tools/proxy/mitm_bench_test.go @@ -206,7 +206,7 @@ func TestMITMCapture_ServerFirst_Fallback(t *testing.T) { t.Fatal(err) } - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() conn, err := dial(ctx, "tcp", tcpServer.Addr().String()) if err != nil { @@ -215,7 +215,7 @@ func TestMITMCapture_ServerFirst_Fallback(t *testing.T) { defer conn.Close() buf := make([]byte, 64) - conn.SetReadDeadline(time.Now().Add(3 * time.Second)) + conn.SetReadDeadline(time.Now().Add(8 * time.Second)) n, err := io.ReadAtLeast(conn, buf, 3) if err != nil { t.Fatalf("expected SSH banner data, got error: %v", err) diff --git a/pkg/tools/register_command.go b/pkg/tools/register_command.go index 0710e330..5842d4ff 100644 --- a/pkg/tools/register_command.go +++ b/pkg/tools/register_command.go @@ -3,12 +3,8 @@ package tools import ( cfg "github.com/chainreactors/aiscan/core/config" "github.com/chainreactors/aiscan/pkg/commands" - gogocmd "github.com/chainreactors/aiscan/pkg/tools/gogo" - neutroncmd "github.com/chainreactors/aiscan/pkg/tools/neutron" "github.com/chainreactors/aiscan/pkg/tools/scan" "github.com/chainreactors/aiscan/pkg/tools/scan/engine" - spraycmd "github.com/chainreactors/aiscan/pkg/tools/spray" - zombiecmd "github.com/chainreactors/aiscan/pkg/tools/zombie" ) func init() { @@ -20,8 +16,6 @@ func init() { if es == nil { return } - logger := deps.GetLogger() - proxy := deps.ScannerProxy var scanOpts []scan.Option for _, o := range deps.ScanOpts { @@ -29,25 +23,13 @@ func init() { scanOpts = append(scanOpts, opt) } } - if proxy != "" { - scanOpts = append(scanOpts, scan.WithProxy(proxy)) + if deps.ScannerProxy != "" { + scanOpts = append(scanOpts, scan.WithProxy(deps.ScannerProxy)) } if es.Gogo != nil && es.Spray != nil { reg.Register(scan.New(es, scanOpts...), "scanner") } - if es.Gogo != nil { - reg.Register(gogocmd.New(es.Gogo).WithLogger(logger).WithProxy(proxy), "scanner") - } - if es.Spray != nil { - reg.Register(spraycmd.New(es.Spray).WithLogger(logger).WithProxy(proxy), "scanner") - } - if es.Zombie != nil { - reg.Register(zombiecmd.New(es.Zombie).WithLogger(logger).WithProxy(proxy), "scanner") - } - if es.Neutron != nil { - reg.Register(neutroncmd.New(es.Neutron, es.Index).WithLogger(logger).WithProxy(proxy), "scanner") - } }, }) } diff --git a/pkg/tools/scan/adapter.go b/pkg/tools/scan/adapter.go index 12ba6cb2..70d2342a 100644 --- a/pkg/tools/scan/adapter.go +++ b/pkg/tools/scan/adapter.go @@ -23,7 +23,7 @@ func (c *Command) runPortDiscoveryCapability(ctx context.Context, discovery disc if target.Ports != "" { ports = target.Ports } - c.logger.Infof("scan capability=%s target=%s ports=%s", capGogoPortscan, target.Target, ports) + c.Logger.Infof("scan capability=%s target=%s ports=%s", capGogoPortscan, target.Target, ports) resultCh, err := engine.GogoScanStream(ctx, c.engines.Gogo, engine.GogoScanOptions{ Target: target.Target, Ports: ports, @@ -31,7 +31,7 @@ func (c *Command) runPortDiscoveryCapability(ctx context.Context, discovery disc Timeout: discovery.Timeout, VersionLevel: discovery.Version, Exploit: discovery.Exploit, - Proxy: c.proxy, + Proxy: c.Proxy, Debug: discovery.Debug, OnStats: func(stats sdktypes.Stats) { emit(statsEvent(capGogoPortscan, stats)) @@ -120,7 +120,7 @@ func (c *Command) runWeakpassCapability(ctx context.Context, flags flags, creden Top: flags.ZombieTop, Users: credentials.Users, Passwords: credentials.Passwords, - Proxy: c.proxy, + Proxy: c.Proxy, Debug: flags.Debug, OnStats: func(stats sdktypes.Stats) { emit(statsEvent(capZombieWeakpass, stats)) diff --git a/pkg/tools/scan/capability.go b/pkg/tools/scan/capability.go index 82f8315b..c188986d 100644 --- a/pkg/tools/scan/capability.go +++ b/pkg/tools/scan/capability.go @@ -96,7 +96,7 @@ func (c *Command) buildCapabilities(flags flags, opts scanOptions, profile profi if !profile.Enabled(name) || !hasSpray(c.engines) { return } - sopts.Proxy = c.proxy + sopts.Proxy = c.Proxy sprayBuilt = true capabilities = append(capabilities, sprayCapability(c, flags, opts.Web, name, sources, sopts, c.runSprayCapability)) } @@ -129,7 +129,7 @@ func (c *Command) buildCapabilities(flags flags, opts scanOptions, profile profi wrapRoutes(acceptsTarget(targetWeb), webSources()...), capWorkers(c.engines.Capacity.Spray, flags.SprayThreads), func(ctx context.Context, e event, emit func(event)) { - c.runSprayCapability(ctx, flags, opts.Web, e.Target, capSprayCrawl, engine.SprayCheckOptions{Crawl: true, CrawlDepth: profile.CrawlDepth, Proxy: c.proxy}, emit) + c.runSprayCapability(ctx, flags, opts.Web, e.Target, capSprayCrawl, engine.SprayCheckOptions{Crawl: true, CrawlDepth: profile.CrawlDepth, Proxy: c.Proxy}, emit) }, )) } @@ -168,13 +168,13 @@ func (c *Command) buildCapabilities(flags flags, opts scanOptions, profile profi } if opts.hasDiscoveryOverrides() && !gogoBuilt { - c.logger.Warnf("scan capability=%s option=port status=ignored reason=engine_unavailable", capGogoPortscan) + c.Logger.Warnf("scan capability=%s option=port status=ignored reason=engine_unavailable", capGogoPortscan) } if opts.hasWebOverrides() && !sprayBuilt { - c.logger.Warnf("scan capability=web_probe option=dict,rule,word,default-dict,advance status=ignored reason=engine_unavailable") + c.Logger.Warnf("scan capability=web_probe option=dict,rule,word,default-dict,advance status=ignored reason=engine_unavailable") } if opts.hasWeakpassOverrides() && !weakpassBuilt { - c.logger.Warnf("scan capability=%s option=user,pwd status=ignored reason=engine_unavailable", capZombieWeakpass) + c.Logger.Warnf("scan capability=%s option=user,pwd status=ignored reason=engine_unavailable", capZombieWeakpass) } for _, builder := range extraCapabilityBuilders { diff --git a/pkg/tools/scan/capability_katana.go b/pkg/tools/scan/capability_katana.go index 41226b0e..f6497b6f 100644 --- a/pkg/tools/scan/capability_katana.go +++ b/pkg/tools/scan/capability_katana.go @@ -122,8 +122,8 @@ func runKatanaCrawl(ctx context.Context, c *Command, e event, depth int, jsMode emit(targetEvent(source, wt.Raw, newWebTarget(wt.Raw, discoveredURL, wt.HostHeader))) }, } - if c.proxy != "" { - options.Proxy = c.proxy + if c.Proxy != "" { + options.Proxy = c.Proxy } gologger.DefaultLogger.SetMaxLevel(levels.LevelSilent) diff --git a/pkg/tools/scan/command.go b/pkg/tools/scan/command.go index 76537806..13816839 100644 --- a/pkg/tools/scan/command.go +++ b/pkg/tools/scan/command.go @@ -19,18 +19,13 @@ import ( ) type Command struct { + toolargs.Base engines *engine.Set parent *agent.Agent deepBrowser DeepBrowserFunc readSkill SkillReader - logger telemetry.Logger - proxy string - workDir string } -func (c *Command) SetWorkDir(dir string) { c.workDir = dir } -func (c *Command) SetProxy(proxy string) { c.proxy = proxy } - type flags struct { Inputs []string `short:"i" long:"input" description:"Input target: URL, IP, IP:port, or CIDR"` ListFile string `short:"l" long:"list" description:"File containing input targets, one per line"` @@ -66,7 +61,8 @@ type flags struct { } func New(engineSet *engine.Set, opts ...Option) *Command { - cmd := &Command{engines: engineSet, logger: telemetry.NopLogger()} + cmd := &Command{engines: engineSet} + cmd.InitLogger(nil) for _, opt := range opts { if opt != nil { opt(cmd) @@ -163,9 +159,9 @@ func (c *Command) execute(ctx context.Context, args []string, stream io.Writer) } if flags.Debug { flags.Trace = true - restoreDebug := telemetry.ActivateDebug(c.logger) + restoreDebug := telemetry.ActivateDebug(c.Logger) defer restoreDebug() - c.logger.Debugf("scan debug enabled") + c.Logger.Debugf("scan debug enabled") } profile, err := profileForMode(flags.Mode) if err != nil { @@ -239,10 +235,10 @@ func (c *Command) execute(ctx context.Context, args []string, stream io.Writer) p.Run(seedsToEvents(seeds)) if c.parent != nil && verifyLevel != "" { - runVerifyPass(ctx, c.parent, c.readSkill, coll, verifyLevel, c.logger) + runVerifyPass(ctx, c.parent, c.readSkill, coll, verifyLevel, c.Logger) } if c.parent != nil && flags.Sniper { - runSniperPass(ctx, c.parent, c.readSkill, coll, c.logger) + runSniperPass(ctx, c.parent, c.readSkill, coll, c.Logger) } coll.Finish() @@ -278,13 +274,13 @@ func (c *Command) execute(ctx context.Context, args []string, stream io.Writer) if flags.OutputFile != "" && !flags.JSON { plainOut := coll.PlainText() if err := writeOutputFile(flags.OutputFile, plainOut); err != nil { - c.logger.Errorf("%s", err.Error()) + c.Logger.Errorf("%s", err.Error()) } } if flags.AssetReportFile != "" { assetOut := coll.AssetReport() if err := writeOutputFile(flags.AssetReportFile, assetOut); err != nil { - c.logger.Errorf("%s", err.Error()) + c.Logger.Errorf("%s", err.Error()) } } return out, coll.StructuredResult(), nil @@ -298,7 +294,7 @@ var scanFileFlags = map[string]bool{ } func (c *Command) resolveRelativePaths(args []string) []string { - return toolargs.ResolveRelativePaths(args, scanFileFlags, c.workDir) + return toolargs.ResolveRelativePaths(args, scanFileFlags, c.WorkDir) } func writeOutputFile(path, content string) error { diff --git a/pkg/tools/scan/options.go b/pkg/tools/scan/options.go index 2cbf0238..aec72bdc 100644 --- a/pkg/tools/scan/options.go +++ b/pkg/tools/scan/options.go @@ -16,15 +16,11 @@ func WithParent(a *agent.Agent) Option { } func WithProxy(proxy string) Option { - return func(c *Command) { c.proxy = proxy } + return func(c *Command) { c.Proxy = proxy } } func WithLogger(logger telemetry.Logger) Option { - return func(c *Command) { - if logger != nil { - c.logger = logger - } - } + return func(c *Command) { c.InitLogger(logger) } } func (c *Command) Configure(opts ...Option) { diff --git a/pkg/tools/spray/register.go b/pkg/tools/spray/register.go new file mode 100644 index 00000000..1a6ba59f --- /dev/null +++ b/pkg/tools/spray/register.go @@ -0,0 +1,22 @@ +package spray + +import ( + "github.com/chainreactors/aiscan/pkg/commands" + "github.com/chainreactors/aiscan/pkg/tools/scan/engine" +) + +func init() { + commands.RegisterFactory(commands.Factory{ + Group: "scanner", + Build: func(deps *commands.Deps, reg *commands.CommandRegistry) { + es, _ := deps.EngineSet.(*engine.Set) + if es == nil || es.Spray == nil { + return + } + reg.Register( + New(es.Spray).WithLogger(deps.GetLogger()).WithProxy(deps.ScannerProxy), + "scanner", + ) + }, + }) +} diff --git a/pkg/tools/spray/spray.go b/pkg/tools/spray/spray.go index 4bbdd513..4a3cbbe6 100644 --- a/pkg/tools/spray/spray.go +++ b/pkg/tools/spray/spray.go @@ -14,32 +14,26 @@ import ( ) type Command struct { - engine *spray.Engine - logger telemetry.Logger - proxy string - workDir string + toolargs.Base + engine *spray.Engine } func New(engine *spray.Engine) *Command { - return &Command{engine: engine, logger: telemetry.NopLogger()} + c := &Command{engine: engine} + c.InitLogger(nil) + return c } func (c *Command) WithLogger(logger telemetry.Logger) *Command { - if logger != nil { - c.logger = logger - } + c.InitLogger(logger) return c } -func (c *Command) SetWorkDir(dir string) { c.workDir = dir } - func (c *Command) WithProxy(proxy string) *Command { - c.proxy = proxy + c.Proxy = proxy return c } -func (c *Command) SetProxy(proxy string) { c.proxy = proxy } - func (c *Command) Name() string { return "spray" } func (c *Command) Usage() string { @@ -51,9 +45,9 @@ func (c *Command) Execute(ctx context.Context, args []string) (err error) { var buf bytes.Buffer debug := toolargs.BoolFlagEnabled(args, "--debug") if debug { - restoreDebug := telemetry.ActivateDebug(c.logger) + restoreDebug := telemetry.ActivateDebug(c.Logger) defer restoreDebug() - c.logger.Debugf("spray debug enabled") + c.Logger.Debugf("spray debug enabled") } if c.engine != nil { c.engine.InstallResourceProvider() @@ -103,13 +97,13 @@ func (c *Command) TestInjectProxy(args []string) []string { } func (c *Command) injectProxy(args []string) []string { - if c.proxy == "" { + if c.Proxy == "" { return args } if toolargs.HasFlag(args, "--proxy") { return args } - return append(args, "--proxy", c.proxy) + return append(args, "--proxy", c.Proxy) } func withDefaultNoBar(args []string) []string { @@ -145,5 +139,5 @@ var sprayFileFlags = map[string]bool{ } func (c *Command) resolveRelativePaths(args []string) []string { - return toolargs.ResolveRelativePaths(args, sprayFileFlags, c.workDir) + return toolargs.ResolveRelativePaths(args, sprayFileFlags, c.WorkDir) } diff --git a/pkg/tools/toolargs/base.go b/pkg/tools/toolargs/base.go new file mode 100644 index 00000000..d27ada65 --- /dev/null +++ b/pkg/tools/toolargs/base.go @@ -0,0 +1,20 @@ +package toolargs + +import "github.com/chainreactors/aiscan/pkg/telemetry" + +type Base struct { + Logger telemetry.Logger + Proxy string + WorkDir string +} + +func (b *Base) SetWorkDir(dir string) { b.WorkDir = dir } +func (b *Base) SetProxy(proxy string) { b.Proxy = proxy } + +func (b *Base) InitLogger(logger telemetry.Logger) { + if logger != nil { + b.Logger = logger + } else { + b.Logger = telemetry.NopLogger() + } +} diff --git a/pkg/tools/zombie/register.go b/pkg/tools/zombie/register.go new file mode 100644 index 00000000..86be9cc9 --- /dev/null +++ b/pkg/tools/zombie/register.go @@ -0,0 +1,22 @@ +package zombie + +import ( + "github.com/chainreactors/aiscan/pkg/commands" + "github.com/chainreactors/aiscan/pkg/tools/scan/engine" +) + +func init() { + commands.RegisterFactory(commands.Factory{ + Group: "scanner", + Build: func(deps *commands.Deps, reg *commands.CommandRegistry) { + es, _ := deps.EngineSet.(*engine.Set) + if es == nil || es.Zombie == nil { + return + } + reg.Register( + New(es.Zombie).WithLogger(deps.GetLogger()).WithProxy(deps.ScannerProxy), + "scanner", + ) + }, + }) +} diff --git a/pkg/tools/zombie/zombie.go b/pkg/tools/zombie/zombie.go index 5e90816f..187a6ba0 100644 --- a/pkg/tools/zombie/zombie.go +++ b/pkg/tools/zombie/zombie.go @@ -13,35 +13,26 @@ import ( ) type Command struct { - engine *sdkzombie.Engine - logger telemetry.Logger - proxy string - workDir string + toolargs.Base + engine *sdkzombie.Engine } func New(engine *sdkzombie.Engine) *Command { - return &Command{engine: engine, logger: telemetry.NopLogger()} + c := &Command{engine: engine} + c.InitLogger(nil) + return c } func (c *Command) WithLogger(logger telemetry.Logger) *Command { - if logger != nil { - c.logger = logger - } + c.InitLogger(logger) return c } -func (c *Command) SetWorkDir(dir string) { c.workDir = dir } - func (c *Command) WithProxy(proxy string) *Command { - c.proxy = proxy + c.Proxy = proxy return c } -// SetProxy stores the proxy URL. Note: the zombie library's RunOptions no -// longer exposes ProxyDial, so proxy is not applied at runtime until upstream -// re-adds support. -func (c *Command) SetProxy(proxy string) { c.proxy = proxy } - func (c *Command) Name() string { return "zombie" } func (c *Command) Usage() string { @@ -52,9 +43,9 @@ func (c *Command) Execute(ctx context.Context, args []string) error { args = c.resolveRelativePaths(args) var buf bytes.Buffer if toolargs.BoolFlagEnabled(args, "--debug") { - restoreDebug := telemetry.ActivateDebug(c.logger) + restoreDebug := telemetry.ActivateDebug(c.Logger) defer restoreDebug() - c.logger.Debugf("zombie debug enabled") + c.Logger.Debugf("zombie debug enabled") } runOpts := zombiecore.RunOptions{Output: &buf} if err := zombiecore.RunWithArgs(ctx, args, runOpts); err != nil { @@ -76,5 +67,5 @@ var zombieFileFlags = map[string]bool{ } func (c *Command) resolveRelativePaths(args []string) []string { - return toolargs.ResolveRelativePaths(args, zombieFileFlags, c.workDir) + return toolargs.ResolveRelativePaths(args, zombieFileFlags, c.WorkDir) } diff --git a/pkg/tui/banner.go b/pkg/tui/banner.go new file mode 100644 index 00000000..6c327cbf --- /dev/null +++ b/pkg/tui/banner.go @@ -0,0 +1,269 @@ +package tui + +import ( + "fmt" + "io" + "os" + "strings" + + cfg "github.com/chainreactors/aiscan/core/config" + outputpkg "github.com/chainreactors/aiscan/core/output" + "golang.org/x/term" +) + +// renderBanner prints a compact welcome block to stderr: title/version, +// resolved model, the session mode, and a short next-step hint. It uses fixed +// ANSI tokens so redirected or recorded sessions do not receive terminal +// background probes. stderr-TTY-only and skipped in quiet mode so redirected +// logs stay clean. Printed once into the scrollback (PTY-forward safe). +func (r *AgentConsole) renderBanner() { + if r.output == nil || r.output.VerbosityLevel() < 0 || r.output.Stderr() == nil { + return + } + if !r.output.tty { + return + } + fmt.Fprint(r.output.Stderr(), r.bannerOutput()) +} + +func (r *AgentConsole) bannerOutput() string { + colorEnabled := r.output != nil && r.output.color.Enabled + provider, model := r.providerModel() + modelText := "not configured - run `aiscan --init`" + modelStyle := ansiWarn + switch { + case provider != "" && model != "": + modelText = provider + " / " + model + modelStyle = ansiAccent + case provider != "": + modelText = provider + modelStyle = ansiAccent + } + + width := r.bannerWidth() + header := ansiTitle("aiscan", colorEnabled) + " " + ansiDim("v"+cfg.Version, colorEnabled) + + var lines []string + lines = append(lines, header) + lines = append(lines, bannerKV("model", modelStyle(modelText, colorEnabled), colorEnabled)) + lines = append(lines, bannerKV("mode", ansiDim(r.sessionSummary(), colorEnabled), colorEnabled)) + lines = append(lines, bannerKV("help", renderInlineCommands([]string{"/help", "/status", "/exit"}, colorEnabled), colorEnabled)) + lines = append(lines, bannerKV("keys", ansiDim("Esc", colorEnabled)+" "+ansiDim("interrupt", colorEnabled)+ansiDim(" ", colorEnabled)+ansiDim("Ctrl+O", colorEnabled)+" "+ansiDim("verbosity", colorEnabled), colorEnabled)) + + box := renderFixedBox(strings.Join(lines, "\n"), width, colorEnabled) + intent := ansiDim("输入目标或任务即可；例如：扫描 192.168.1.10 的 Web 风险", colorEnabled) + + var b strings.Builder + fmt.Fprintln(&b) + fmt.Fprintln(&b, box) + fmt.Fprintln(&b, " "+intent) + if notice := strings.TrimSpace(r.startupNotice); notice != "" { + fmt.Fprintln(&b, " "+ansiWarn("⚠ "+notice, colorEnabled)) + } + fmt.Fprintln(&b) + return b.String() +} + +func (r *AgentConsole) bannerWidth() int { + const ( + minWidth = 44 + defaultWidth = 64 + maxWidth = 78 + ) + width := defaultWidth + if r != nil && r.terminal != nil && r.terminal.Control != nil { + if columns, _ := r.terminal.Control.Size(); columns > 0 { + width = columns - 4 + } + } else if r != nil && r.output != nil && r.output.Stderr() != nil { + if columns := writerTerminalWidth(r.output.Stderr()); columns > 0 { + width = columns - 4 + } + } + if width < minWidth { + return minWidth + } + if width > maxWidth { + return maxWidth + } + return width +} + +func writerTerminalWidth(w io.Writer) int { + file, ok := w.(*os.File) + if !ok { + return 0 + } + width, _, err := term.GetSize(int(file.Fd())) + if err != nil || width <= 0 { + return 0 + } + return width +} + +func bannerKV(label, value string, colorEnabled bool) string { + return ansiDim(fmt.Sprintf("%-9s", label), colorEnabled) + value +} + +func renderFixedBox(body string, width int, colorEnabled bool) string { + const minInnerWidth = 16 + innerWidth := width - 4 + if innerWidth < minInnerWidth { + innerWidth = minInnerWidth + } + lines := strings.Split(body, "\n") + for _, line := range lines { + if n := visibleRuneLen(line); n > innerWidth { + innerWidth = n + } + } + + border := func(s string) string { return ansiDim(s, colorEnabled) } + var b strings.Builder + fmt.Fprintf(&b, "%s\n", border("╭"+strings.Repeat("─", innerWidth+2)+"╮")) + for _, line := range lines { + padding := innerWidth - visibleRuneLen(line) + if padding < 0 { + padding = 0 + } + fmt.Fprintf(&b, "%s %s%s %s\n", + border("│"), + line, + strings.Repeat(" ", padding), + border("│")) + } + fmt.Fprint(&b, border("╰"+strings.Repeat("─", innerWidth+2)+"╯")) + return b.String() +} + +func visibleRuneLen(s string) int { + return len([]rune(outputpkg.StripANSI(s))) +} + +func ansiWrap(s, code string, enabled bool) string { + return outputpkg.NewColor(enabled).Wrap(s, code) +} + +func ansiTitle(s string, enabled bool) string { + return ansiWrap(s, outputpkg.ANSIBold+outputpkg.ANSICyan, enabled) +} + +func ansiAccent(s string, enabled bool) string { + return ansiWrap(s, outputpkg.ANSICyan, enabled) +} + +func ansiWarn(s string, enabled bool) string { + return ansiWrap(s, outputpkg.ANSIYellow, enabled) +} + +func ansiDim(s string, enabled bool) string { + return ansiWrap(s, outputpkg.ANSIDim, enabled) +} + +func renderInlineCommands(commands []string, colorEnabled bool) string { + parts := make([]string, 0, len(commands)) + for _, command := range commands { + parts = append(parts, ansiAccent(command, colorEnabled)) + } + return strings.Join(parts, ansiDim(" ", colorEnabled)) +} + +func (r *AgentConsole) sessionSummary() string { + var parts []string + if r != nil && r.output != nil { + switch r.output.mode { + case ModeStatic: + parts = append(parts, "static") + case ModeForwarded: + parts = append(parts, "forwarded") + default: + parts = append(parts, "pty") + } + if r.output.stream.enabled { + parts = append(parts, "stream") + } else if r.output.Markdown() { + parts = append(parts, "pretty") + } else { + parts = append(parts, "plain") + } + } + if r != nil && r.option != nil { + if space := strings.TrimSpace(r.option.Space); space != "" { + parts = append(parts, "space "+space) + } + } + if len(parts) == 0 { + return "pty" + } + return strings.Join(parts, " · ") +} + +func (r *AgentConsole) providerModel() (string, string) { + if r.appInfo.Commands == nil { + return "", "" + } + pc := r.appInfo.ProviderConfig + return pc.Provider, pc.Model +} + +func (r *AgentConsole) renderHelp() string { + colorEnabled := r.output != nil && r.output.color.Enabled + cmds := r.allCommands() + rows := make([]helpRow, 0, len(cmds)+3) + for _, c := range cmds { + if c.Hidden { + continue + } + rows = append(rows, helpRow{Command: c.Name, Detail: c.Description}) + } + rows = append(rows, helpRow{}) + rows = append(rows, helpRow{Command: "普通文本", Detail: "直接发送自然语言任务"}) + rows = append(rows, helpRow{Command: "! <命令>", Detail: "直接执行 bash/伪命令（跳过 LLM）"}) + return r.renderPanel("commands", renderHelpRows(rows, colorEnabled), colorEnabled) +} + +func (r *AgentConsole) renderStatus() string { + colorEnabled := r.output != nil && r.output.color.Enabled + info := CollectStatus(r.replSession(), r.sessionSummary(), agentConsoleHistoryPath()) + rows := []helpRow{ + {Command: "model", Detail: info.Provider + " / " + info.Model}, + {Command: "render", Detail: info.Mode}, + {Command: "task", Detail: info.Task}, + {Command: "ioa", Detail: info.IOA}, + {Command: "history", Detail: info.History}, + } + if info.Skills != "" { + rows = append(rows, helpRow{Command: "skills", Detail: info.Skills}) + } + return r.renderPanel("status", renderHelpRows(rows, colorEnabled), colorEnabled) +} + +type helpRow struct { + Command string + Detail string +} + +const helpRowCommandWidth = 18 + +func renderHelpRows(rows []helpRow, colorEnabled bool) string { + var b strings.Builder + for _, row := range rows { + if row.Command == "" && row.Detail == "" { + b.WriteByte('\n') + continue + } + command := ansiAccent(fmt.Sprintf("%-*s", helpRowCommandWidth, row.Command), colorEnabled) + detail := ansiDim(row.Detail, colorEnabled) + fmt.Fprintf(&b, "%s%s\n", command, detail) + } + return strings.TrimRight(b.String(), "\n") +} + +func (r *AgentConsole) renderPanel(title, body string, colorEnabled bool) string { + title = strings.TrimSpace(title) + if title == "" { + title = "aiscan" + } + header := ansiTitle(title, colorEnabled) + return "\n" + renderFixedBox(header+"\n"+body, r.bannerWidth(), colorEnabled) + "\n\n" +} diff --git a/pkg/tui/console.go b/pkg/tui/console.go index 51b2515f..56186172 100644 --- a/pkg/tui/console.go +++ b/pkg/tui/console.go @@ -8,7 +8,6 @@ import ( "io" "os" "path/filepath" - "sort" "strings" "sync" "sync/atomic" @@ -21,10 +20,8 @@ import ( "github.com/chainreactors/aiscan/pkg/agent" ioaclient "github.com/chainreactors/ioa/client" "github.com/reeflective/console" - "github.com/reeflective/readline/inputrc" rlterm "github.com/reeflective/readline/terminal" "github.com/spf13/cobra" - "golang.org/x/term" ) const agentPromptCommandName = "__prompt" @@ -134,225 +131,6 @@ func NewAgentConsoleWithTerminal(ctx context.Context, option *cfg.Option, appInf return repl } -func configureAgentReadline(c *console.Console) { - if c == nil { - return - } - shell := c.Shell() - cfg := shell.Config - _ = cfg.Set("autocomplete", true) - _ = cfg.Set("usage-hint-always", false) - _ = cfg.Set("history-autosuggest", true) - _ = cfg.Set("show-all-if-ambiguous", true) - _ = cfg.Set("show-all-if-unmodified", true) - _ = cfg.Set("menu-complete-display-prefix", true) - _ = cfg.Set("page-completions", false) - _ = cfg.Set("completion-query-items", 1000) - _ = cfg.Set("bell-style", "none") - _ = cfg.Set("enable-bracketed-paste", false) - // Bind Tab to menu-complete so arrow keys navigate the dropdown. - for _, keymap := range []string{"emacs", "emacs-standard", "vi-insert"} { - _ = cfg.Bind(keymap, `\t`, "menu-complete", false) - _ = cfg.Bind(keymap, inputrc.Unescape(`\e[Z`), "menu-complete-backward", false) - } -} - -func (r *AgentConsole) configureInterruptKey() { - if r == nil || r.console == nil || r.console.Shell() == nil { - return - } - shell := r.console.Shell() - shell.Keymap.Register(map[string]func(){ - agentConsoleInterruptCommandName: func() { - r.handleEscapeInterruptKey() - }, - }) - escape := inputrc.Unescape(`\e`) - for _, keymap := range []string{"emacs", "emacs-standard"} { - _ = shell.Config.Bind(keymap, escape, agentConsoleInterruptCommandName, false) - } -} - -func (r *AgentConsole) configureCtrlCKey() { - if r == nil || r.console == nil || r.console.Shell() == nil { - return - } - shell := r.console.Shell() - shell.Keymap.Register(map[string]func(){ - agentConsoleCtrlCCommandName: func() { - r.handleCtrlC() - }, - }) - ctrlC := inputrc.Unescape(`\C-c`) - for _, keymap := range []string{"emacs", "emacs-standard"} { - _ = shell.Config.Bind(keymap, ctrlC, agentConsoleCtrlCCommandName, false) - } -} - -func (r *AgentConsole) handleCtrlC() { - if r.InterruptCurrentRun() { - return - } - if r.pendingExit.Load() { - os.Exit(0) - } - r.pendingExit.Store(true) - fmt.Fprintf(r.stderr, " Press Ctrl+C again to exit\n") - go func() { - time.Sleep(3 * time.Second) - r.pendingExit.Store(false) - }() - shell := r.console.Shell() - shell.Display.AcceptLine() - shell.History.Accept(false, false, errors.New(os.Interrupt.String())) -} - -func (r *AgentConsole) configureVerbosityToggleKey() { - if r == nil || r.console == nil || r.console.Shell() == nil { - return - } - shell := r.console.Shell() - shell.Keymap.Register(map[string]func(){ - agentConsoleToggleVerbosityCommandName: func() { - r.handleToggleVerbosity() - }, - }) - ctrlO := inputrc.Unescape(`\C-o`) - for _, keymap := range []string{"emacs", "emacs-standard"} { - _ = shell.Config.Bind(keymap, ctrlO, agentConsoleToggleVerbosityCommandName, false) - } -} - -func (r *AgentConsole) handleToggleVerbosity() { - out := r.ensureOutput() - if out == nil { - return - } - current := out.VerbosityLevel() - next := (current + 1) % 3 - out.SetVerbosity(next) - label := out.VerbosityLabel() - if out.color.Enabled { - fmt.Fprintf(r.stderr, "\n%s %s\n", - out.dim("verbosity:"), - out.color.Wrap(label, outputpkg.ANSICyan)) - } else { - fmt.Fprintf(r.stderr, "\nverbosity: %s\n", label) - } -} - -func (r *AgentConsole) handleEscapeInterruptKey() { - if r == nil || r.console == nil || r.console.Shell() == nil { - return - } - shell := r.console.Shell() - pending := string(shell.Keys.Read()) - if pending == "" { - pending = readPendingTerminalBytes(agentConsoleEscapeSequenceWait) - } - keymap := string(shell.Keymap.Main()) - if feed, ok := agentConsoleEscapeSequenceFeed(shell.Config.Binds[keymap], pending); ok { - shell.Keys.Feed(true, []rune(feed)...) - return - } - if pending != "" { - shell.Keys.Feed(true, []rune(pending)...) - return - } - r.InterruptCurrentRun() -} - -func agentConsoleEscapeSequenceFeed(binds map[string]inputrc.Bind, pending string) (string, bool) { - if len(binds) == 0 || pending == "" { - return "", false - } - sequence := inputrc.Unescape(`\e`) + pending - matches := make([]string, 0, 4) - for seq := range binds { - readlineSeq := agentConsoleReadlineSequence(seq) - if len(readlineSeq) <= 1 || !strings.HasPrefix(readlineSeq, inputrc.Unescape(`\e`)) { - continue - } - if strings.HasPrefix(sequence, readlineSeq) { - matches = append(matches, seq) - } - } - sort.Slice(matches, func(i, j int) bool { - left := agentConsoleReadlineSequence(matches[i]) - right := agentConsoleReadlineSequence(matches[j]) - if len(left) == len(right) { - return left < right - } - return len(left) > len(right) - }) - for _, seq := range matches { - bind := binds[seq] - replacement, ok := agentConsoleEquivalentNonEscapeBind(binds, bind) - if !ok { - continue - } - return replacement + sequence[len(agentConsoleReadlineSequence(seq)):], true - } - return "", false -} - -func agentConsoleEquivalentNonEscapeBind(binds map[string]inputrc.Bind, target inputrc.Bind) (string, bool) { - if target.Action == "" { - return "", false - } - candidates := make([]string, 0, 4) - for seq, bind := range binds { - if bind.Action != target.Action || bind.Macro != target.Macro || strings.HasPrefix(agentConsoleReadlineSequence(seq), inputrc.Unescape(`\e`)) { - continue - } - candidates = append(candidates, seq) - } - sort.Slice(candidates, func(i, j int) bool { - if len(candidates[i]) == len(candidates[j]) { - return candidates[i] < candidates[j] - } - return len(candidates[i]) < len(candidates[j]) - }) - if len(candidates) == 0 { - return agentConsoleFallbackNonEscapeBind(target) - } - return candidates[0], true -} - -func agentConsoleFallbackNonEscapeBind(target inputrc.Bind) (string, bool) { - switch target.Action { - case "previous-history", "history-search-backward": - return inputrc.Unescape(`\C-p`), true - case "next-history", "history-search-forward": - return inputrc.Unescape(`\C-n`), true - case "backward-char", "vi-backward-char": - return inputrc.Unescape(`\C-b`), true - case "forward-char", "vi-forward-char": - return inputrc.Unescape(`\C-f`), true - case "beginning-of-line": - return inputrc.Unescape(`\C-a`), true - case "end-of-line": - return inputrc.Unescape(`\C-e`), true - default: - return "", false - } -} - -func agentConsoleReadlineSequence(seq string) string { - if seq == "" { - return "" - } - converted := make([]rune, 0, len(seq)) - for _, r := range seq { - if inputrc.IsMeta(r) { - converted = append(converted, inputrc.Esc, inputrc.Demeta(r)) - continue - } - converted = append(converted, r) - } - return string(converted) -} - func (r *AgentConsole) Start() error { r.renderBanner() defer r.stopController() @@ -514,201 +292,6 @@ func (r *AgentConsole) executeArgs(ctx context.Context, args []string) error { return root.Execute() } -// renderBanner prints a compact welcome block to stderr: title/version, -// resolved model, the session mode, and a short next-step hint. It uses fixed -// ANSI tokens so redirected or recorded sessions do not receive terminal -// background probes. stderr-TTY-only and skipped in quiet mode so redirected -// logs stay clean. Printed once into the scrollback (PTY-forward safe). -func (r *AgentConsole) renderBanner() { - if r.output == nil || r.output.VerbosityLevel() < 0 || r.output.Stderr() == nil { - return - } - if !r.output.tty { - return - } - fmt.Fprint(r.output.Stderr(), r.bannerOutput()) -} - -func (r *AgentConsole) bannerOutput() string { - colorEnabled := r.output != nil && r.output.color.Enabled - provider, model := r.providerModel() - modelText := "not configured - run `aiscan --init`" - modelStyle := ansiWarn - switch { - case provider != "" && model != "": - modelText = provider + " / " + model - modelStyle = ansiAccent - case provider != "": - modelText = provider - modelStyle = ansiAccent - } - - width := r.bannerWidth() - header := ansiTitle("aiscan", colorEnabled) + " " + ansiDim("v"+cfg.Version, colorEnabled) - - var lines []string - lines = append(lines, header) - lines = append(lines, bannerKV("model", modelStyle(modelText, colorEnabled), colorEnabled)) - lines = append(lines, bannerKV("mode", ansiDim(r.sessionSummary(), colorEnabled), colorEnabled)) - lines = append(lines, bannerKV("help", renderInlineCommands([]string{"/help", "/status", "/exit"}, colorEnabled), colorEnabled)) - lines = append(lines, bannerKV("keys", ansiDim("Esc", colorEnabled)+" "+ansiDim("interrupt", colorEnabled)+ansiDim(" ", colorEnabled)+ansiDim("Ctrl+O", colorEnabled)+" "+ansiDim("verbosity", colorEnabled), colorEnabled)) - - box := renderFixedBox(strings.Join(lines, "\n"), width, colorEnabled) - intent := ansiDim("输入目标或任务即可；例如：扫描 192.168.1.10 的 Web 风险", colorEnabled) - - var b strings.Builder - fmt.Fprintln(&b) - fmt.Fprintln(&b, box) - fmt.Fprintln(&b, " "+intent) - if notice := strings.TrimSpace(r.startupNotice); notice != "" { - fmt.Fprintln(&b, " "+ansiWarn("⚠ "+notice, colorEnabled)) - } - fmt.Fprintln(&b) - return b.String() -} - -func (r *AgentConsole) bannerWidth() int { - const ( - minWidth = 44 - defaultWidth = 64 - maxWidth = 78 - ) - width := defaultWidth - if r != nil && r.terminal != nil && r.terminal.Control != nil { - if columns, _ := r.terminal.Control.Size(); columns > 0 { - width = columns - 4 - } - } else if r != nil && r.output != nil && r.output.Stderr() != nil { - if columns := writerTerminalWidth(r.output.Stderr()); columns > 0 { - width = columns - 4 - } - } - if width < minWidth { - return minWidth - } - if width > maxWidth { - return maxWidth - } - return width -} - -func writerTerminalWidth(w io.Writer) int { - file, ok := w.(*os.File) - if !ok { - return 0 - } - width, _, err := term.GetSize(int(file.Fd())) - if err != nil || width <= 0 { - return 0 - } - return width -} - -func bannerKV(label, value string, colorEnabled bool) string { - return ansiDim(fmt.Sprintf("%-9s", label), colorEnabled) + value -} - -func renderFixedBox(body string, width int, colorEnabled bool) string { - const minInnerWidth = 16 - innerWidth := width - 4 - if innerWidth < minInnerWidth { - innerWidth = minInnerWidth - } - lines := strings.Split(body, "\n") - for _, line := range lines { - if n := visibleRuneLen(line); n > innerWidth { - innerWidth = n - } - } - - border := func(s string) string { return ansiDim(s, colorEnabled) } - var b strings.Builder - fmt.Fprintf(&b, "%s\n", border("╭"+strings.Repeat("─", innerWidth+2)+"╮")) - for _, line := range lines { - padding := innerWidth - visibleRuneLen(line) - if padding < 0 { - padding = 0 - } - fmt.Fprintf(&b, "%s %s%s %s\n", - border("│"), - line, - strings.Repeat(" ", padding), - border("│")) - } - fmt.Fprint(&b, border("╰"+strings.Repeat("─", innerWidth+2)+"╯")) - return b.String() -} - -func visibleRuneLen(s string) int { - return len([]rune(outputpkg.StripANSI(s))) -} - -func ansiWrap(s, code string, enabled bool) string { - return outputpkg.NewColor(enabled).Wrap(s, code) -} - -func ansiTitle(s string, enabled bool) string { - return ansiWrap(s, outputpkg.ANSIBold+outputpkg.ANSICyan, enabled) -} - -func ansiAccent(s string, enabled bool) string { - return ansiWrap(s, outputpkg.ANSICyan, enabled) -} - -func ansiWarn(s string, enabled bool) string { - return ansiWrap(s, outputpkg.ANSIYellow, enabled) -} - -func ansiDim(s string, enabled bool) string { - return ansiWrap(s, outputpkg.ANSIDim, enabled) -} - -func renderInlineCommands(commands []string, colorEnabled bool) string { - parts := make([]string, 0, len(commands)) - for _, command := range commands { - parts = append(parts, ansiAccent(command, colorEnabled)) - } - return strings.Join(parts, ansiDim(" ", colorEnabled)) -} - -func (r *AgentConsole) sessionSummary() string { - var parts []string - if r != nil && r.output != nil { - switch r.output.mode { - case ModeStatic: - parts = append(parts, "static") - case ModeForwarded: - parts = append(parts, "forwarded") - default: - parts = append(parts, "pty") - } - if r.output.stream.enabled { - parts = append(parts, "stream") - } else if r.output.Markdown() { - parts = append(parts, "pretty") - } else { - parts = append(parts, "plain") - } - } - if r != nil && r.option != nil { - if space := strings.TrimSpace(r.option.Space); space != "" { - parts = append(parts, "space "+space) - } - } - if len(parts) == 0 { - return "pty" - } - return strings.Join(parts, " · ") -} - -func (r *AgentConsole) providerModel() (string, string) { - if r.appInfo.Commands == nil { - return "", "" - } - pc := r.appInfo.ProviderConfig - return pc.Provider, pc.Model -} - func (r *AgentConsole) replSession() *Session { s := &Session{ Ctx: r.ctx, @@ -985,72 +568,6 @@ func wrapCommand(cmd Command, s *Session) *cobra.Command { return cc } -// --------------------------------------------------------------------------- -// Rendering: help, status, panels -// --------------------------------------------------------------------------- - -func (r *AgentConsole) renderHelp() string { - colorEnabled := r.output != nil && r.output.color.Enabled - cmds := r.allCommands() - rows := make([]helpRow, 0, len(cmds)+3) - for _, c := range cmds { - if c.Hidden { - continue - } - rows = append(rows, helpRow{Command: c.Name, Detail: c.Description}) - } - rows = append(rows, helpRow{}) - rows = append(rows, helpRow{Command: "普通文本", Detail: "直接发送自然语言任务"}) - rows = append(rows, helpRow{Command: "! <命令>", Detail: "直接执行 bash/伪命令（跳过 LLM）"}) - return r.renderPanel("commands", renderHelpRows(rows, colorEnabled), colorEnabled) -} - -func (r *AgentConsole) renderStatus() string { - colorEnabled := r.output != nil && r.output.color.Enabled - info := CollectStatus(r.replSession(), r.sessionSummary(), agentConsoleHistoryPath()) - rows := []helpRow{ - {Command: "model", Detail: info.Provider + " / " + info.Model}, - {Command: "render", Detail: info.Mode}, - {Command: "task", Detail: info.Task}, - {Command: "ioa", Detail: info.IOA}, - {Command: "history", Detail: info.History}, - } - if info.Skills != "" { - rows = append(rows, helpRow{Command: "skills", Detail: info.Skills}) - } - return r.renderPanel("status", renderHelpRows(rows, colorEnabled), colorEnabled) -} - -type helpRow struct { - Command string - Detail string -} - -const helpRowCommandWidth = 18 - -func renderHelpRows(rows []helpRow, colorEnabled bool) string { - var b strings.Builder - for _, row := range rows { - if row.Command == "" && row.Detail == "" { - b.WriteByte('\n') - continue - } - command := ansiAccent(fmt.Sprintf("%-*s", helpRowCommandWidth, row.Command), colorEnabled) - detail := ansiDim(row.Detail, colorEnabled) - fmt.Fprintf(&b, "%s%s\n", command, detail) - } - return strings.TrimRight(b.String(), "\n") -} - -func (r *AgentConsole) renderPanel(title, body string, colorEnabled bool) string { - title = strings.TrimSpace(title) - if title == "" { - title = "aiscan" - } - header := ansiTitle(title, colorEnabled) - return "\n" + renderFixedBox(header+"\n"+body, r.bannerWidth(), colorEnabled) + "\n\n" -} - func (r *AgentConsole) ensureOutput() *AgentOutput { if r.output == nil { r.output = NewAgentOutput(r.option) diff --git a/pkg/tui/format.go b/pkg/tui/format.go index e26d57de..05ba4054 100644 --- a/pkg/tui/format.go +++ b/pkg/tui/format.go @@ -246,9 +246,10 @@ func summarizeChatMessage(msg agent.ChatMessage) (role string, contentLen int, t // --------------------------------------------------------------------------- var ( - agentMarkdownRenderer *glamour.TermRenderer - agentMarkdownRendererErr error - agentMarkdownRendererOnce sync.Once + agentMarkdownRenderer *glamour.TermRenderer + agentMarkdownRendererErr error + agentMarkdownRendererW int + agentMarkdownRendererMu sync.Mutex ) func renderAgentMarkdown(content string, enabled bool) string { @@ -275,20 +276,22 @@ func renderAgentMarkdown(content string, enabled bool) string { } func getAgentMarkdownRenderer() (*glamour.TermRenderer, error) { - agentMarkdownRendererOnce.Do(func() { - opts := []glamour.TermRendererOption{ - glamour.WithAutoStyle(), - // Auto-detect the richest profile the terminal advertises (truecolor - // -> 256 -> ANSI) instead of pinning 16-color ANSI, so markdown answers - // render with real depth on modern terminals. - glamour.WithColorProfile(termenv.ColorProfile()), - glamour.WithEmoji(), - } - if w := terminalWidth(); w > 0 { - opts = append(opts, glamour.WithWordWrap(w)) - } - agentMarkdownRenderer, agentMarkdownRendererErr = glamour.NewTermRenderer(opts...) - }) + w := terminalWidth() + agentMarkdownRendererMu.Lock() + defer agentMarkdownRendererMu.Unlock() + if agentMarkdownRenderer != nil && w == agentMarkdownRendererW { + return agentMarkdownRenderer, agentMarkdownRendererErr + } + opts := []glamour.TermRendererOption{ + glamour.WithAutoStyle(), + glamour.WithColorProfile(termenv.ColorProfile()), + glamour.WithEmoji(), + } + if w > 0 { + opts = append(opts, glamour.WithWordWrap(w)) + } + agentMarkdownRenderer, agentMarkdownRendererErr = glamour.NewTermRenderer(opts...) + agentMarkdownRendererW = w return agentMarkdownRenderer, agentMarkdownRendererErr } diff --git a/pkg/tui/keybindings.go b/pkg/tui/keybindings.go new file mode 100644 index 00000000..6108ea02 --- /dev/null +++ b/pkg/tui/keybindings.go @@ -0,0 +1,233 @@ +package tui + +import ( + "errors" + "fmt" + "os" + "sort" + "strings" + "time" + + outputpkg "github.com/chainreactors/aiscan/core/output" + "github.com/reeflective/console" + "github.com/reeflective/readline/inputrc" +) + +func configureAgentReadline(c *console.Console) { + if c == nil { + return + } + shell := c.Shell() + cfg := shell.Config + _ = cfg.Set("autocomplete", true) + _ = cfg.Set("usage-hint-always", false) + _ = cfg.Set("history-autosuggest", true) + _ = cfg.Set("show-all-if-ambiguous", true) + _ = cfg.Set("show-all-if-unmodified", true) + _ = cfg.Set("menu-complete-display-prefix", true) + _ = cfg.Set("page-completions", false) + _ = cfg.Set("completion-query-items", 1000) + _ = cfg.Set("bell-style", "none") + _ = cfg.Set("enable-bracketed-paste", false) + // Bind Tab to menu-complete so arrow keys navigate the dropdown. + for _, keymap := range []string{"emacs", "emacs-standard", "vi-insert"} { + _ = cfg.Bind(keymap, `\t`, "menu-complete", false) + _ = cfg.Bind(keymap, inputrc.Unescape(`\e[Z`), "menu-complete-backward", false) + } +} + +func (r *AgentConsole) configureInterruptKey() { + if r == nil || r.console == nil || r.console.Shell() == nil { + return + } + shell := r.console.Shell() + shell.Keymap.Register(map[string]func(){ + agentConsoleInterruptCommandName: func() { + r.handleEscapeInterruptKey() + }, + }) + escape := inputrc.Unescape(`\e`) + for _, keymap := range []string{"emacs", "emacs-standard"} { + _ = shell.Config.Bind(keymap, escape, agentConsoleInterruptCommandName, false) + } +} + +func (r *AgentConsole) configureCtrlCKey() { + if r == nil || r.console == nil || r.console.Shell() == nil { + return + } + shell := r.console.Shell() + shell.Keymap.Register(map[string]func(){ + agentConsoleCtrlCCommandName: func() { + r.handleCtrlC() + }, + }) + ctrlC := inputrc.Unescape(`\C-c`) + for _, keymap := range []string{"emacs", "emacs-standard"} { + _ = shell.Config.Bind(keymap, ctrlC, agentConsoleCtrlCCommandName, false) + } +} + +func (r *AgentConsole) handleCtrlC() { + if r.InterruptCurrentRun() { + return + } + if r.pendingExit.Load() { + os.Exit(0) + } + r.pendingExit.Store(true) + fmt.Fprintf(r.stderr, " Press Ctrl+C again to exit\n") + go func() { + time.Sleep(3 * time.Second) + r.pendingExit.Store(false) + }() + shell := r.console.Shell() + shell.Display.AcceptLine() + shell.History.Accept(false, false, errors.New(os.Interrupt.String())) +} + +func (r *AgentConsole) configureVerbosityToggleKey() { + if r == nil || r.console == nil || r.console.Shell() == nil { + return + } + shell := r.console.Shell() + shell.Keymap.Register(map[string]func(){ + agentConsoleToggleVerbosityCommandName: func() { + r.handleToggleVerbosity() + }, + }) + ctrlO := inputrc.Unescape(`\C-o`) + for _, keymap := range []string{"emacs", "emacs-standard"} { + _ = shell.Config.Bind(keymap, ctrlO, agentConsoleToggleVerbosityCommandName, false) + } +} + +func (r *AgentConsole) handleToggleVerbosity() { + out := r.ensureOutput() + if out == nil { + return + } + current := out.VerbosityLevel() + next := (current + 1) % 3 + out.SetVerbosity(next) + label := out.VerbosityLabel() + if out.color.Enabled { + fmt.Fprintf(r.stderr, "\n%s %s\n", + out.dim("verbosity:"), + out.color.Wrap(label, outputpkg.ANSICyan)) + } else { + fmt.Fprintf(r.stderr, "\nverbosity: %s\n", label) + } +} + +func (r *AgentConsole) handleEscapeInterruptKey() { + if r == nil || r.console == nil || r.console.Shell() == nil { + return + } + shell := r.console.Shell() + pending := string(shell.Keys.Read()) + if pending == "" { + pending = readPendingTerminalBytes(agentConsoleEscapeSequenceWait) + } + keymap := string(shell.Keymap.Main()) + if feed, ok := agentConsoleEscapeSequenceFeed(shell.Config.Binds[keymap], pending); ok { + shell.Keys.Feed(true, []rune(feed)...) + return + } + if pending != "" { + shell.Keys.Feed(true, []rune(pending)...) + return + } + r.InterruptCurrentRun() +} + +func agentConsoleEscapeSequenceFeed(binds map[string]inputrc.Bind, pending string) (string, bool) { + if len(binds) == 0 || pending == "" { + return "", false + } + sequence := inputrc.Unescape(`\e`) + pending + matches := make([]string, 0, 4) + for seq := range binds { + readlineSeq := agentConsoleReadlineSequence(seq) + if len(readlineSeq) <= 1 || !strings.HasPrefix(readlineSeq, inputrc.Unescape(`\e`)) { + continue + } + if strings.HasPrefix(sequence, readlineSeq) { + matches = append(matches, seq) + } + } + sort.Slice(matches, func(i, j int) bool { + left := agentConsoleReadlineSequence(matches[i]) + right := agentConsoleReadlineSequence(matches[j]) + if len(left) == len(right) { + return left < right + } + return len(left) > len(right) + }) + for _, seq := range matches { + bind := binds[seq] + replacement, ok := agentConsoleEquivalentNonEscapeBind(binds, bind) + if !ok { + continue + } + return replacement + sequence[len(agentConsoleReadlineSequence(seq)):], true + } + return "", false +} + +func agentConsoleEquivalentNonEscapeBind(binds map[string]inputrc.Bind, target inputrc.Bind) (string, bool) { + if target.Action == "" { + return "", false + } + candidates := make([]string, 0, 4) + for seq, bind := range binds { + if bind.Action != target.Action || bind.Macro != target.Macro || strings.HasPrefix(agentConsoleReadlineSequence(seq), inputrc.Unescape(`\e`)) { + continue + } + candidates = append(candidates, seq) + } + sort.Slice(candidates, func(i, j int) bool { + if len(candidates[i]) == len(candidates[j]) { + return candidates[i] < candidates[j] + } + return len(candidates[i]) < len(candidates[j]) + }) + if len(candidates) == 0 { + return agentConsoleFallbackNonEscapeBind(target) + } + return candidates[0], true +} + +func agentConsoleFallbackNonEscapeBind(target inputrc.Bind) (string, bool) { + switch target.Action { + case "previous-history", "history-search-backward": + return inputrc.Unescape(`\C-p`), true + case "next-history", "history-search-forward": + return inputrc.Unescape(`\C-n`), true + case "backward-char", "vi-backward-char": + return inputrc.Unescape(`\C-b`), true + case "forward-char", "vi-forward-char": + return inputrc.Unescape(`\C-f`), true + case "beginning-of-line": + return inputrc.Unescape(`\C-a`), true + case "end-of-line": + return inputrc.Unescape(`\C-e`), true + default: + return "", false + } +} + +func agentConsoleReadlineSequence(seq string) string { + if seq == "" { + return "" + } + converted := make([]rune, 0, len(seq)) + for _, r := range seq { + if inputrc.IsMeta(r) { + converted = append(converted, inputrc.Esc, inputrc.Demeta(r)) + continue + } + converted = append(converted, r) + } + return string(converted) +} diff --git a/pkg/util/format.go b/pkg/util/format.go index 19db0057..62c02492 100644 --- a/pkg/util/format.go +++ b/pkg/util/format.go @@ -17,6 +17,9 @@ func FormatSize(bytes int) string { } func FormatNumber(n int) string { + if n < 0 { + return "-" + FormatNumber(-n) + } if n < 1000 { return fmt.Sprintf("%d", n) } diff --git a/web/frontend/src/lib/scan-result.ts b/web/frontend/src/lib/scan-result.ts index b34b1e7a..a7168e78 100644 --- a/web/frontend/src/lib/scan-result.ts +++ b/web/frontend/src/lib/scan-result.ts @@ -682,7 +682,7 @@ function normalizePriority(priority?: string): FindingPriority { function itemToPriority(item: AssetItem): FindingPriority { const p = (item.status || '').toLowerCase() - if (p === 'confirmed' || p === 'critical') return 'critical' + if (p === 'critical') return 'critical' if (p === 'high') return 'high' if (p === 'medium' || p === 'inconclusive') return 'medium' if (p === 'low') return 'low'