使用 argcomplete 为 impacket/netexec 等工具添加命令行自动补全

在日常的渗透测试与安全研究中，我们经常与功能复杂的工具集打交道。以 Impacket 为例，一个软件包内包含了 60 多个工具（如 secretsdump.py、psexec.py、GetUserSPNs.py 等），每个工具又有数十乃至上百个参数。依赖死记硬背不现实，频繁敲击 --help 又会打断心流。

Shell 自带的文件补全仅能补路径，对参数补全无能为力。我们真正需要的是这样流畅的体验：

$ secretsdump.py -<TAB>
-h                    --help                --just-dc
-no-pass              --just-dc-user        --just-dhcp
--just-dns            --k                ...

这种体验依赖于 动态补全 机制：当按下 Tab 键时，Shell 将当前命令行输入交由程序本体处理，程序根据其内部的 argparse 定义实时计算并返回补全候选。这正是 argcomplete 的核心价值所在。

---

一、配置篇：两种路径，一套解法

我们将分别以 NetExec（自动识别） 和 Impacket（手动定制） 为例，展示两种典型的配置场景。

1.1 体验零配置：NetExec 的自动识别

NetExec 的 pip 安装版，在安装后即可享受补全功能，用户无需任何额外配置。这背后的功臣是 argcomplete 的 全局激活 机制。

• 前置条件：确保已安装 argcomplete 并激活全局补全。

  pip install argcomplete
  activate-global-python-argcomplete  # 安装全局钩子
  # 根据提示，重启 Shell 或 source 配置文件

• 效果：在新终端中输入 netexec -<TAB>，补全列表将自动出现。

这种“零配置”的便利性，源于全局钩子对命令的 静态检查（详见原理部分）。它自动识别出 NetExec 是一个符合规范的 console_script，且其模块内包含了 PYTHON_ARGCOMPLETE_OK 标记。

1.2 全量手工定制：为 62 个 Impacket 工具注入补全灵魂

Impacket 的 apt 版脚本（位于 /usr/share/doc/python3-impacket/examples/）并未包含任何补全支持。我们需要对其进行一次性的、幂等的“手术”。

第一步：复制脚本至用户目录 (PATH 优先) 为避免污染系统源包，同时确保我们的修改版本优先被调用，将脚本复制到 ~/.local/bin。

cp /usr/share/doc/python3-impacket/examples/*.py ~/.local/bin/

第二步：使用增强脚本进行幂等插桩

以下是完整的 addAutocomplete.py 脚本，它对目录下所有 .py 文件执行幂等的 argcomplete 插桩，并在修改后自动进行语法校验：

#!/usr/bin/env python3
"""addAutocomplete.py - 幂等增强版

对目录下所有 .py 脚本插桩 argcomplete:
1. 在 import argparse 行后插入 `import argcomplete`（幂等：已有则跳过）
2. 在 `if len(sys.argv) == 1:` 前插入 `argcomplete.autocomplete(parser)`（同缩进）
3. 兜底：无 sys.argv 判断时，插在 `parser.parse_args()` 前
4. 修改后自动执行 py_compile 校验语法
"""
import os
import sys
import re
import py_compile


def modify_file(file_path):
    """修改单个 Python 文件，插入 argcomplete 支持"""
    with open(file_path, 'r', encoding='utf-8') as f:
        lines = f.readlines()

    modified = False

    # 幂等保护：已有 argcomplete import 则跳过该文件
    if any('import argcomplete' in line for line in lines):
        print(f"SKIP (already patched): {file_path}")
        return False

    # 1. 在第一个 import argparse 行后插入 import argcomplete
    insert_idx = None
    for i, line in enumerate(lines):
        if re.match(r'^\s*import argparse\b', line) or re.match(r'^\s*from argparse\b', line):
            insert_idx = i + 1
            break
    if insert_idx is not None:
        lines.insert(insert_idx, 'import argcomplete\n')
        modified = True

    # 2. 优先: if len(sys.argv) == 1: 前插 autocomplete(parser)
    pattern_if = re.compile(r'^(\s*)if len\(sys\.argv\) ?== ?1:')
    pattern_parse_args = re.compile(r'^(\s*).*parser\.parse_args\(\)')

    autocomplete_inserted = False
    for i, line in enumerate(lines):
        match_if = pattern_if.match(line)
        if match_if:
            indent = match_if.group(1)
            lines.insert(i, f'{indent}argcomplete.autocomplete(parser)\n')
            modified = True
            autocomplete_inserted = True
            break

    # 3. 兜底: parser.parse_args() 前插
    if not autocomplete_inserted:
        for i, line in enumerate(lines):
            match_parse_args = pattern_parse_args.match(line)
            if match_parse_args:
                indent = match_parse_args.group(1)
                lines.insert(i, f'{indent}argcomplete.autocomplete(parser)\n')
                modified = True
                break

    if modified:
        with open(file_path, 'w', encoding='utf-8') as f:
            f.writelines(lines)
        try:
            py_compile.compile(file_path, doraise=True)
            print(f"MODIFIED + COMPILE OK: {file_path}")
        except py_compile.PyCompileError as e:
            print(f"MODIFIED but COMPILE FAILED: {file_path}: {e}")
            return False
    else:
        print(f"NO CHANGE (pattern not found): {file_path}")

    return modified


def main(directory):
    """遍历目录下所有 .py 文件并修改"""
    if not os.path.isdir(directory):
        print(f"Error: {directory} is not a valid directory.")
        sys.exit(1)

    count = 0
    for filename in sorted(os.listdir(directory)):
        if filename.endswith('.py'):
            file_path = os.path.join(directory, filename)
            if modify_file(file_path):
                count += 1

    print(f"\n{count} files modified.")


if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: ./addAutocomplete.py <directory>")
        sys.exit(1)
    main(sys.argv[1])

运行脚本：

python3 addAutocomplete.py ~/.local/bin/
# MODIFIED + COMPILE OK: ~/.local/bin/atexec.py
# MODIFIED + COMPILE OK: ~/.local/bin/secretsdump.py
# ... 62 files modified.

每个文件改完都跑 py_compile 校验语法，确保插桩不破坏脚本。

第三步：生成并注册 Shell 补全函数 argcomplete 提供了 register-python-argcomplete 命令来生成注册代码。我们可以利用它为所有工具批量生成 compdef 条目，并追加到 ~/.zshrc 中。

# 1. 生成一个工具的注册模板，提取核心函数
register-python-argcomplete psexec.py > /tmp/register_template.txt

# 2. 将模板中的 __python_argcomplete_run、_python_argcomplete 等函数，
#    以及所有 62 个工具的 compdef 行，整理后追加到 .zshrc

关键配置片段 (~/.zshrc):

# ===== impacket argcomplete 补全 (62 tools) =====
# ... [核心函数 __python_argcomplete_run, _python_argcomplete] ...
compdef _python_argcomplete secretsdump.py
compdef _python_argcomplete psexec.py
# ... (其余 60 个工具)
# ===== end impacket argcomplete =====

第四步：解决“补全冲突” (常见于 Oh-My-Zsh) 若使用了 fzf 等插件，它可能劫持 Tab 键。必须在 .zshrc 加载 Oh-My-Zsh 之前 设置环境变量，将 Tab 行为归还给补全系统：

export fzf_default_completion=expand-or-complete
source $ZSH/oh-my-zsh.sh

重启终端后，psexec.py -<TAB> 将展示带描述的补全列表，大功告成。

---

二、原理篇：argcomplete 的工作流与两大激活策略

理解了“怎么做”，我们深入“为什么”，掌握 argcomplete 的设计哲学。

2.1 核心矛盾与暴力解法

• 矛盾：argparse 的参数定义在程序运行时才构建，而补全需要提前知道参数列表。
• 传统方案：为每个工具手写补全脚本，维护成本高，易过时。
• argcomplete 的解法：补全即运行。在补全时直接启动程序本体，让程序自身根据实时的 parser 对象生成候选。程序即定义，永不脱节。

2.2 架构：一次完整的补全握手

整个交互过程由 Shell 端和 Python 端协作完成。

┌────────────┐  1. 用户按 Tab          ┌─────────────────────────┐
│  Shell     │ ──────────────────────▶ │  Python 脚本 (带插桩)   │
│  (zsh/bash)│  设置环境变量, 运行脚本  │                         │
│            │ ◀────────────────────── │  3. 通过 fd 8 / 临时文件 │
└────────────┘  2. 返回补全候选 (VT分隔)│  返回 argcomplete 候选  │

关键环境变量协议：

• _ARGCOMPLETE：触发标志 (1/2/3 表示不同的脚本执行方式)。
• COMP_LINE / COMP_POINT：完整的命令行文本与光标位置。Python 端据此重建 argv。
• _ARGCOMPLETE_SHELL：指定 Shell 类型 (zsh/bash)，用于适配输出格式。

输出通道 (fd 8)： 补全时，脚本的标准输出 (stdout) 和错误 (stderr) 均被重定向至 /dev/null，防止干扰终端。真正的补全数据通过 文件描述符 8 (fd 8) 传递，若无 fd 8，则回退至临时文件模式。

插桩点 autocomplete(parser)： 这是脚本中被注入的核心函数。它在 parse_args() 之前被调用：

• 无 _ARGCOMPLETE 环境变量：立即返回，业务逻辑零开销。
• 有 _ARGCOMPLETE 环境变量：解析环境变量，生成补全候选，输出后调用 sys.exit(0)，绝不会进入真正的攻击/业务逻辑。

2.3 两种激活策略：显式注册与全局兜底

argcomplete 提供了两种激活方式，其机制截然不同。

• 显式注册 (Explicit Registration)

  • 命令：register-python-argcomplete <脚本>
  • 机制：为指定的命令生成并绑定一个专用的补全函数 (complete -F _python_argcomplete mytool 或 compdef _python_argcomplete mytool)。
  • 特点：精准、快速、确定性高。每增加一个新工具，需要重新注册一次。

• 全局激活 (Global Activation)

  • 命令：activate-global-python-argcomplete
  • 机制：安装一个兜底钩子。
    • Bash：注册为默认补全器 (complete -D)。
    • Zsh：注册为默认补全器 (#compdef -default-)。
  • 特点：为所有没有显式注册的命令提供补全尝试。关键：它与显式注册共存，当一个命令拥有显式注册时，全局钩子不会被调用。

2.4 全局激活的“自动识别”是静态的

全局钩子并非“运行每个程序来探测”，而是执行纯静态检查，这解释了为何 NetExec 免配置而 Impacket 不行。

_check_console_script 的检查逻辑如下：

1. 脚本的 basename 必须命中一个已安装的 console_scripts 入口点。
2. 脚本内容必须匹配 console_script 的模板 (如 sys.exit(function()))。
3. 对应的模块 __init__.py 头部必须包含 PYTHON_ARGCOMPLETE_OK 标记。

实测对比：

# NetExec: 满足全部条件, 检查通过
$ python3 -m argcomplete._check_console_script /usr/bin/netexec; echo $?
0

# Impacket 手动脚本: 非 console_script, 检查失败
$ python3 -m argcomplete._check_console_script ~/.local/bin/secretsdump.py; echo $?
1

结论：NetExec 的零配置，是作者主动在代码中植入了标记。而像 Impacket 这类手动维护的脚本集，必须依赖显式注册。

---

三、排错实录：两个典型的陷阱

Bug 1：Tab 被 fzf-completion 劫持

• 现象：按下 Tab 无反应或触发目录跳转。
• 根因：fzf 插件绑定了 Tab，其 fallback 机制覆盖了 Zsh 的原生补全系统 (_comps)，导致显式注册被忽略。
• 修复：在加载 Oh-My-Zsh 之前，设置 export fzf_default_completion=expand-or-complete。

Bug 2：注册块中的空行导致补全为空

• 现象：_comps 注册正确，但补全结果为空。
• 根因：从模板复制代码时，\ 续行符之间若插入空行，Zsh 会将其解析为多个独立的命令，导致环境变量 (COMP_LINE 等) 无法正确导出，Python 端抛出 KeyError 而崩溃。
• 修复：确保注册块中的命令替换（completions=( ... )）为单行，或保证续行符连续，无多余空行。

---

四、总结与最佳实践

1. 核心认知：argcomplete 通过“补全即运行”的理念，利用环境变量协议，将补全逻辑与程序定义合二为一，从根本上解决了参数维护的难题。
2. 激活方式：
   • 显式注册 适合自研工具或手动管理的脚本集，精准可靠。
   • 全局激活 为已集成补全标记的工具（如 NetExec）提供“无感”体验，作为兜底方案。
3. 自动识别本质：全局钩子的“自动”是有条件的，它依赖于静态的 console_script 结构和代码内的 PYTHON_ARGCOMPLETE_OK 标记。理解这一点，就能准确判断何时该用何种方式。
4. 检查清单：
   • 插桩后使用 py_compile 验证语法。
   • 注册块中的 Shell 代码避免多余空行，尤其是续行符附近。
   • 若使用 fzf 等插件，确保其补全模式配置正确，且加载顺序无误。
   • 修改后，务必在新终端中测试，因为补全注册发生在 Shell 启动时。

原文 https://blog.csdn.net/2301_79518550/article/details/156302360