main.py

#!/usr/bin/env python3
"""
protobuf_rev - Protobuf 逆向工具
从去符号的二进制文件中提取 protobuf 定义，支持:
  - C (protobuf-c)
  - C++ (Google protobuf)
  - Go (google.golang.org/protobuf)
  - Rust (prost / rust-protobuf)
  - 其他嵌入了序列化 FileDescriptorProto 的程序

用法:
  python main.py <binary> [-o output_dir] [-s scanner] [--protoc PATH] [--compile]
"""

import argparse
import os
import sys
import shutil
from pathlib import Path

from binary_reader import BinaryReader
from proto_generator import ProtoFileGenerator, fdproto_to_proto_string
from scanner_protobuf_c import (
    scan_protobuf_c, scan_protobuf_c_enums, protobuf_c_to_proto,
    ProtobufCType, ProtobufCLabel,
)
from scanner_cpp import scan_cpp_protobuf
from scanner_go import scan_go_protobuf, scan_go_protobuf_raw
from scanner_raw import scan_raw_descriptors


# ANSI colors
class C:
    RESET = '\033[0m'
    BOLD = '\033[1m'
    RED = '\033[91m'
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    BLUE = '\033[94m'
    MAGENTA = '\033[95m'
    CYAN = '\033[96m'
    DIM = '\033[2m'


BANNER = f"""{C.CYAN}{C.BOLD}
  ╔═══════════════════════════════════════════════╗
  ║         protobuf_rev - Protobuf 逆向工具      ║
  ║     C / C++ / Go / Rust  ➜  .proto 文件       ║
  ╚═══════════════════════════════════════════════╝{C.RESET}
"""


def print_info(msg):
    print(f"  {C.BLUE}[*]{C.RESET} {msg}")

def print_ok(msg):
    print(f"  {C.GREEN}[+]{C.RESET} {msg}")

def print_warn(msg):
    print(f"  {C.YELLOW}[!]{C.RESET} {msg}")

def print_err(msg):
    print(f"  {C.RED}[-]{C.RESET} {msg}")

def print_section(title):
    print(f"\n  {C.BOLD}{C.MAGENTA}{'─' * 50}")
    print(f"  {title}")
    print(f"  {'─' * 50}{C.RESET}")


def run_scan(filepath: str, scanners: list = None, output_dir: str = '.', protoc_path: str = None, compile_proto: bool = False):
    """执行扫描并输出结果"""
    print(BANNER)

    # 加载二进制文件
    print_info(f"加载文件: {C.BOLD}{filepath}{C.RESET}")
    reader = BinaryReader(filepath)
    print_info(f"格式: {reader.info}")

    all_scanners = scanners or ['protobuf-c', 'cpp', 'go', 'raw']
    all_proto_contents = []
    all_fd_protos = []

    # ═══════════════ protobuf-c 扫描 ═══════════════
    if 'protobuf-c' in all_scanners:
        print_section("🔍 扫描 protobuf-c (C 语言)")
        messages = scan_protobuf_c(reader)
        enums = scan_protobuf_c_enums(reader)

        if messages:
            print_ok(f"发现 {C.BOLD}{len(messages)}{C.RESET} 个 ProtobufCMessageDescriptor")
            for msg in messages:
                print(f"      {C.CYAN}├─{C.RESET} {msg.name} ({msg.n_fields} fields)")
                for field in msg.fields:
                    label = ProtobufCLabel.NAMES.get(field.label, '')
                    type_name = ProtobufCType.NAMES.get(field.type, '?')
                    print(f"      {C.DIM}│  {label} {type_name} {field.name} = {field.id}{C.RESET}")

            if enums:
                print_ok(f"发现 {len(enums)} 个 ProtobufCEnumDescriptor")
                for e in enums:
                    print(f"      {C.CYAN}├─{C.RESET} {e.name} ({len(e.values)} values)")

            proto_content = protobuf_c_to_proto(messages, enums)
            if proto_content:
                all_proto_contents.append(('protobuf_c', proto_content))
        else:
            print_warn("未发现 protobuf-c 结构体")

    # ═══════════════ C++ protobuf 扫描 ═══════════════
    if 'cpp' in all_scanners:
        print_section("🔍 扫描 C++ protobuf (序列化 FileDescriptorProto)")
        fd_protos = scan_cpp_protobuf(reader)

        if fd_protos:
            print_ok(f"发现 {C.BOLD}{len(fd_protos)}{C.RESET} 个 FileDescriptorProto")
            for fd in fd_protos:
                all_fd_protos.append(fd)
                print(f"      {C.CYAN}├─{C.RESET} {fd.name} (package: {fd.package or '<none>'})")
                for msg in fd.message_type:
                    print(f"      {C.DIM}│  message {msg.name} ({len(msg.field)} fields){C.RESET}")
                for enum in fd.enum_type:
                    print(f"      {C.DIM}│  enum {enum.name} ({len(enum.value)} values){C.RESET}")
        else:
            print_warn("未发现 C++ protobuf 序列化描述符")

    # ═══════════════ Go protobuf 扫描 ═══════════════
    if 'go' in all_scanners:
        existing_names = {fd.name for fd in all_fd_protos}
        print_section("🔍 扫描 Go protobuf (gzip 压缩描述符)")
        fd_protos = scan_go_protobuf(reader)

        # 如果 gzip 方法没找到，尝试 raw 方法
        if not fd_protos:
            fd_protos = scan_go_protobuf_raw(reader)

        # 去重: 跳过已由 C++ 扫描器发现的
        fd_protos = [fd for fd in fd_protos if fd.name not in existing_names]

        if fd_protos:
            print_ok(f"发现 {C.BOLD}{len(fd_protos)}{C.RESET} 个 FileDescriptorProto (Go)")
            for fd in fd_protos:
                all_fd_protos.append(fd)
                print(f"      {C.CYAN}├─{C.RESET} {fd.name} (package: {fd.package or '<none>'})")
                for msg in fd.message_type:
                    print(f"      {C.DIM}│  message {msg.name} ({len(msg.field)} fields){C.RESET}")
                for enum in fd.enum_type:
                    print(f"      {C.DIM}│  enum {enum.name} ({len(enum.value)} values){C.RESET}")
        else:
            print_warn("未发现 Go protobuf 描述符")

    # ═══════════════ 通用扫描 ═══════════════
    if 'raw' in all_scanners:
        # 避免重复: 跳过已经找到的
        existing_names = {fd.name for fd in all_fd_protos}

        print_section("🔍 通用扫描 (Rust/其他)")
        fd_protos = scan_raw_descriptors(reader)
        new_protos = [fd for fd in fd_protos if fd.name not in existing_names]

        if new_protos:
            print_ok(f"发现 {C.BOLD}{len(new_protos)}{C.RESET} 个额外的 FileDescriptorProto")
            for fd in new_protos:
                all_fd_protos.append(fd)
                print(f"      {C.CYAN}├─{C.RESET} {fd.name} (package: {fd.package or '<none>'})")
                for msg in fd.message_type:
                    print(f"      {C.DIM}│  message {msg.name} ({len(msg.field)} fields){C.RESET}")
        else:
            print_warn("未发现额外的 protobuf 描述符")

    # ═══════════════ 生成 .proto 文件 ═══════════════
    print_section("📝 生成 .proto 文件")

    os.makedirs(output_dir, exist_ok=True)
    generated_files = []

    # 从 FileDescriptorProto 生成
    for fd in all_fd_protos:
        proto_content = fdproto_to_proto_string(fd)
        # 使用原始文件名或生成一个
        out_name = fd.name if fd.name else 'unknown.proto'
        # 移除路径前缀
        out_name = os.path.basename(out_name)
        out_path = os.path.join(output_dir, out_name)

        with open(out_path, 'w', encoding='utf-8') as f:
            f.write(proto_content)
        generated_files.append(out_path)
        print_ok(f"生成: {C.BOLD}{out_path}{C.RESET}")
        # 打印内容预览
        for line in proto_content.strip().split('\n'):
            print(f"      {C.DIM}{line}{C.RESET}")

    # 从 protobuf-c 生成
    for tag, proto_content in all_proto_contents:
        out_path = os.path.join(output_dir, 'protobuf_c_extracted.proto')
        with open(out_path, 'w', encoding='utf-8') as f:
            f.write(proto_content)
        generated_files.append(out_path)
        print_ok(f"生成: {C.BOLD}{out_path}{C.RESET}")
        for line in proto_content.strip().split('\n'):
            print(f"      {C.DIM}{line}{C.RESET}")

    if not generated_files:
        print_err("未能提取到任何 protobuf 定义")
        return 1

    # ═══════════════ 可选: 编译 .proto ═══════════════
    if compile_proto:
        print_section("⚙️  编译 .proto 文件")
        protoc = protoc_path or shutil.which('protoc')
        if not protoc:
            print_warn("未找到 protoc, 跳过编译")
        else:
            for proto_file in generated_files:
                proto_dir = os.path.dirname(os.path.abspath(proto_file))
                cmd = f'"{protoc}" --python_out="{proto_dir}" --proto_path="{proto_dir}" "{os.path.abspath(proto_file)}"'
                print_info(f"运行: {cmd}")
                ret = os.system(cmd)
                if ret == 0:
                    print_ok(f"编译成功: {proto_file}")
                else:
                    print_err(f"编译失败 (返回码 {ret}): {proto_file}")

    # 汇总
    print_section("📊 汇总")
    total_msgs = sum(len(fd.message_type) for fd in all_fd_protos)
    total_msgs += sum(len(msgs) for _, msgs in [(t, c) for t, c in all_proto_contents] if isinstance(msgs, list)) if False else 0
    print_ok(f"共提取 {C.BOLD}{len(generated_files)}{C.RESET} 个 .proto 文件")
    print_ok(f"FileDescriptorProto: {len(all_fd_protos)} 个")
    if all_proto_contents:
        print_ok(f"protobuf-c 结构体: {len(all_proto_contents)} 个")
    print()

    return 0


def main():
    parser = argparse.ArgumentParser(
        description='protobuf_rev - 从二进制文件中提取 protobuf 定义',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  python main.py ./pwn                          # 自动检测并提取
  python main.py ./pwn -o output/               # 指定输出目录
  python main.py ./pwn -s protobuf-c            # 只扫描 protobuf-c
  python main.py ./pwn -s cpp -s go             # 扫描 C++ 和 Go
  python main.py ./pwn --compile                # 提取并编译为 _pb2.py
  python main.py ./pwn --protoc /path/to/protoc # 指定 protoc 路径
        """,
    )
    parser.add_argument('binary', help='目标二进制文件路径')
    parser.add_argument('-o', '--output', default='.', help='输出目录 (默认: 当前目录)')
    parser.add_argument('-s', '--scanner', action='append', choices=['protobuf-c', 'cpp', 'go', 'raw', 'all'],
                        help='指定扫描器 (可多选, 默认: all)')
    parser.add_argument('--protoc', default=None, help='protoc 路径')
    parser.add_argument('--compile', action='store_true', help='生成后自动编译 .proto 文件')

    args = parser.parse_args()

    if not os.path.isfile(args.binary):
        print_err(f"文件不存在: {args.binary}")
        return 1

    scanners = args.scanner
    if not scanners or 'all' in scanners:
        scanners = ['protobuf-c', 'cpp', 'go', 'raw']

    return run_scan(
        filepath=args.binary,
        scanners=scanners,
        output_dir=args.output,
        protoc_path=args.protoc,
        compile_proto=args.compile,
    )


if __name__ == '__main__':
    sys.exit(main())