#!/usr/bin/env python3
"""Batch decompiler module - CLI entry point is in cli.py."""

import os
import sys
import argparse
import subprocess
import multiprocessing
from pathlib import Path
from typing import List, Optional, Tuple

# IDA Pro 要处理的文件类型
FILE_TYPES = {
    'executable': {
        'description': '可执行文件 (Mach-O, ELF, PE)',
        'extensions': ['', '.out', '.exe', '.elf', '.bin'],
        'magic_bytes': {
            b'\x7fELF': 'ELF',
            b'\xCF\xFA\xED\xFE': 'Mach-O 64-bit',  # 0xFEEDFACE reversed
            b'\xCE\xFA\xED\xFE': 'Mach-O 32-bit',
            b'\x4D\x5A': 'PE/DOS executable',
        }
    },
    'sharedlib': {
        'description': '共享库/动态库 (Mach-O dylib, ELF so, DLL)',
        'extensions': ['.so', '.dylib', '.dll'],
        'magic_bytes': {
            b'\x7fELF': 'ELF shared library',
            b'\xCF\xFA\xED\xFE': 'Mach-O 64-bit dylib',
            b'\xCE\xFA\xED\xFE': 'Mach-O 32-bit dylib',
            b'\x4D\x5A': 'PE DLL',
        }
    }
}


def detect_file_type(file_path: Path) -> Optional[Tuple[str, str]]:
    """
    检测文件类型

    Args:
        file_path: 文件路径

    Returns:
        (file_category, file_description) 或 None
    """
    try:
        with open(file_path, 'rb') as f:
            # 读取前 4 字节用于 magic number 检测
            magic = f.read(4)

            for file_type, config in FILE_TYPES.items():
                for magic_bytes, desc in config['magic_bytes'].items():
                    if magic.startswith(magic_bytes):
                        return (file_type, desc)

            # 对于 Mach-O，还需要检查 32/64 位
            if len(magic) >= 4:
                # Mach-O 字节序检查
                if magic[:2] in (b'\xfe\xed', b'\xce\xfa'):
                    # 读取更多字节判断
                    f.seek(0)
                    macho_header = f.read(20)
                    if len(macho_header) >= 20:
                        # 检查 cpu type
                        import struct
                        if len(magic) >= 4:
                            # 简单检查 - 64 位 Mach-O 的 cputype 位 0x01000007
                            pass

            # 扩展名检查
            for file_type, config in FILE_TYPES.items():
                if any(file_path.suffix.lower() == ext for ext in config['extensions']):
                    # 只有当扩展名匹配但 magic 不确定时才返回
                    # 如果已经通过 magic 检测到，就不需要这里
                    pass

    except (IOError, PermissionError):
        pass

    return None


def is_binary_file(file_path: Path) -> bool:
    """判断文件是否为可执行文件或共享库"""
    try:
        # 排除已知的非二进制文件
        excluded_extensions = {'.py', '.pyc', '.txt', '.md', '.json', '.xml', '.yaml',
                               '.yml', '.log', '.cfg', '.ini', '.conf', '.html', '.css',
                               '.js', '.ts', '.java', '.c', '.cpp', '.h', '.hpp'}

        if file_path.suffix.lower() in excluded_extensions:
            return False

        # 检查文件大小
        if file_path.stat().st_size < 64:  # 太小的文件不是二进制
            return False

        # 检查是否是文本文件
        try:
            with open(file_path, 'rb') as f:
                chunk = f.read(8192)
                # 检查空字节 - 二进制文件通常包含
                if b'\x00' in chunk:
                    return True
                # 检查是否是纯文本
                try:
                    chunk.decode('utf-8')
                    # 纯文本，可能是脚本，检查 shebang
                    if chunk.startswith(b'#!'):
                        return True  # 脚本也可能需要处理
                except UnicodeDecodeError:
                    return True
        except (IOError, PermissionError):
            return False

        # 检测 magic number
        with open(file_path, 'rb') as f:
            magic = f.read(4)

        for config in FILE_TYPES.values():
            if magic[:2] in (b'\x7fE', b'\xCF\xFA', b'\xCE\xFA', b'\x4D\x5A'):
                return True

        # 检查扩展名
        for config in FILE_TYPES.values():
            if any(file_path.suffix.lower() == ext for ext in config['extensions']):
                return True

    except (IOError, OSError, PermissionError):
        pass

    return False


def find_binaries(directory: Path) -> List[Path]:
    """在目录中递归查找所有二进制文件"""
    binaries = []

    for root, dirs, files in os.walk(directory):
        # 跳过常见非二进制目录
        dirs[:] = [d for d in dirs if d not in {'.git', '.svn', '.idea', '.vscode',
                                                 '__pycache__', 'node_modules',
                                                 '.cache', 'build', 'dist'}]

        for filename in files:
            file_path = Path(root) / filename

            # 快速初筛
            if not is_binary_file(file_path):
                continue

            # 详细检测
            file_type = detect_file_type(file_path)
            if file_type:
                binaries.append(file_path)
                print(f"  Found: {file_type[0]:15} -> {file_path}")

    return binaries


def create_idapython_script(output_dir: Path) -> Path:
    """
    创建 IDAPython 脚本用于导出反编译的 C 代码

    Args:
        output_dir: 输出目录

    Returns:
        脚本文件路径
    """
    script_content = """
import os
import idaapi
import idc
import idautils

# 获取输出目录
output_dir = r\"{output_dir}\"

# 获取二进制文件名（不带扩展名）
input_file = idc.get_input_file_path()
bin_name = os.path.splitext(os.path.basename(input_file))[0]
output_file = os.path.join(output_dir, bin_name + ".c")

print(f"[IDAPython] 开始导出反编译代码: {output_file}")

# 确保输出目录存在
os.makedirs(output_dir, exist_ok=True)

# 获取所有函数
functions = list(idautils.Functions())
print(f"[IDAPython] 找到 {len(functions)} 个函数")

# 生成 C 代码
try:
    # 使用 IDAPython 的 decompile 接口
    import ida_hexrays

    if ida_hexrays.init_hexrays_plugin(0):
        print("[IDAPython] Hex-Rays 反编译器可用")
    else:
        print("[IDAPython] Hex-Rays 反编译器不可用，尝试使用内置反编译")

except ImportError:
    print("[IDAPython] Hex-Rays 模块不可用")

# 导出 C 代码
# 方法1: 使用 idc.Export() 如果可用
try:
    if idaapi.export_to_c(output_file, 0, idc.BADADDR, 0) == 0:
        print(f"[IDAPython] 成功导出 C 代码: {output_file}")
    else:
        print(f"[IDAPython] 导出失败")
except Exception as e:
    print(f"[IDAPython] 导出异常: {e}")
    # 方法2: 尝试使用 idc.Save Base
    try:
        # 生成伪代码
        idc.generate_file(idc.GFL_CDECOMP, output_file, 0, idc.BADADDR, 0)
        print(f"[IDAPython] 使用 GFL_CDECOMP 导出: {output_file}")
    except Exception as e2:
        print(f"[IDAPython] 两个方法都失败: {e2}")

# 退出 IDA
print("[IDAPython] 完成，退出 IDA...")
idaapi.qexit(0)
""".format(output_dir=output_dir)

    script_path = Path(__file__).parent / "idapython_export_c.py"
    script_path.write_text(script_content)
    return script_path


def create_idapython_script_v2(output_dir: Path) -> Path:
    """
    创建改进版 IDAPython 脚本，使用 IDC 的 generate_file
    """
    script_content = f"""import os
import idaapi
import idc

# 输出目录
output_dir = r"{output_dir}"

# 获取输入文件信息
input_file = idc.get_input_file_path()
bin_name = os.path.splitext(os.path.basename(input_file))[0]
output_file = os.path.join(output_dir, bin_name + ".c")

print(f"Processing: {input_file}")
print(f"Output: {output_file}")

# 确保输出目录存在
os.makedirs(output_dir, exist_ok=True)

# 使用 IDC 的 generate_file 函数导出 C 代码
# GFL_CDECOMP = 0x0004 (C source file)
try:
    result = idc.generate_file(idc.GFL_CDECOMP, output_file, 0, idc.BADADDR, 0)
    if result == 0:
        print(f"SUCCESS: Generated {output_file}")
    else:
        print(f"FAILED: generate_file returned {result}")
except Exception as e:
    print(f"ERROR: {e}")

# 保存数据库（可选）
# idc.save_database(output_file.replace('.c', '.i64'))

# 退出
idaapi.qexit(0)
"""

    script_path = Path(__file__).parent / "idapython_export_c.py"
    script_path.write_text(script_content)
    return script_path


def decompile_binary(
    binary_path: Path,
    output_dir: Path,
    idat_path: Path,
    timeout: int = 300,
    verbose: bool = False
) -> Tuple[bool, str]:
    """
    使用 IDA Pro decompile 一个二进制文件

    Args:
        binary_path: 二进制文件路径
        output_dir: 输出目录
        idat_path: idat 可执行文件路径
        timeout: 超时时间(秒)
        verbose: 是否显示详细输出

    Returns:
        (success, message)
    """
    bin_name = binary_path.stem
    output_file = output_dir / f"{bin_name}.c"

    # 创建临时脚本
    script_content = f"""import os
import idaapi
import idc

# 输出目录
output_dir = r"{output_dir}"

# 获取输入文件信息
input_file = idc.get_input_file_path()
bin_name = os.path.splitext(os.path.basename(input_file))[0]
output_file = os.path.join(output_dir, bin_name + ".c")

print(f"Processing: {{input_file}}")
print(f"Output: {{output_file}}")

# 确保输出目录存在
os.makedirs(output_dir, exist_ok=True)

# 使用 IDC 的 generate_file 函数导出 C 代码
# GFL_CDECOMP = 0x0004 (C source file)
try:
    result = idc.generate_file(idc.GFL_CDECOMP, output_file, 0, idc.BADADDR, 0)
    if result == 0:
        print(f"SUCCESS: Generated {{output_file}}")
    else:
        print(f"FAILED: generate_file returned {{result}}")
except Exception as e:
    print(f"ERROR: {{e}}")

# 退出
idaapi.qexit(0)
"""

    # 写入脚本文件
    script_path = Path(tempfile.gettempdir()) / f"idapython_EXPORT_C_{bin_name}.py"
    script_path.write_text(script_content)

    # 构建 idat 命令
    # -A: 自动模式
    # -S: 执行脚本
    # -L: 日志文件
    idat_cmd = [
        str(idat_path),
        "-A",
        f"-S{idapython_SCRIPT_PATH}",
        f"-L{output_dir / f'{bin_name}.log'}",
        str(binary_path)
    ]

    if verbose:
        print(f"  Command: {' '.join(idat_cmd)}")

    try:
        result = subprocess.run(
            idat_cmd,
            timeout=timeout,
            capture_output=True,
            text=True
        )

        if result.returncode == 0 and output_file.exists():
            return True, f"Successfully generated {output_file}"
        else:
            error_msg = result.stderr.strip() if result.stderr else "Unknown error"
            return False, f"Failed: {error_msg}"

    except subprocess.TimeoutExpired:
        return False, f"Timeout after {timeout}s"
    except Exception as e:
        return False, f"Exception: {e}"
    finally:
        # 清理临时脚本
        try:
            script_path.unlink()
        except:
            pass


def decompile_binary_simple(
    binary_path: Path,
    output_dir: Path,
    idat_path: Path,
    timeout: int = 300,
    verbose: bool = False
) -> Tuple[bool, str]:
    """
    使用 IDA Pro decompile 一个二进制文件 (简化版)

    使用 -Cm 参数让 IDA 自动加载并处理
    """
    bin_name = binary_path.stem
    output_file = output_dir / f"{bin_name}.c"

    # 创建脚本文件
    script_content = f"""import idc
import idaapi
import os

output_dir = r"{output_dir}"
bin_name = os.path.splitext(idc.get_input_file_path())[0]
output_file = os.path.join(output_dir, bin_name + ".c")

os.makedirs(output_dir, exist_ok=True)

# 导出 C 代码
result = idc.generate_file(idc.GFL_CDECOMP, output_file, 0, idc.BADADDR, 0)
if result == 0:
    print(f"SUCCESS: {{output_file}}")
else:
    print(f"FAILED: {{result}}")

idaapi.qexit(0)
"""

    # 生成唯一的脚本文件名
    import time
    script_name = f"idapython_export_{int(time.time() * 1000)}.py"
    script_path = Path(tempfile.gettempdir()) / script_name
    script_path.write_text(script_content)

    try:
        # 构建命令
        # Linux/macOS: idat
        # Windows: idat64 或 idat32
        cmd = [
            str(idat_path),
            "-A",       # Batch mode
            f"-S{script_path}",  # Run script
            str(binary_path)
        ]

        if verbose:
            print(f"  Running: {' '.join(cmd)}")

        result = subprocess.run(
            cmd,
            timeout=timeout,
            capture_output=True,
            text=True,
            env={**os.environ, "IDA_LOG_FILE": str(output_dir / f"{bin_name}.log")}
        )

        if output_file.exists():
            return True, f"Generated {output_file}"
        else:
            stderr = result.stderr.strip() if result.stderr else "No error output"
            return False, f"Output file not created: {stderr}"

    except subprocess.TimeoutExpired:
        return False, f"Timeout after {timeout}s"
    except Exception as e:
        return False, f"Exception: {e}"
    finally:
        try:
            script_path.unlink(missing_ok=True)
        except:
            pass


def decompile_binary_v3(
    binary_path: Path,
    output_dir: Path,
    idat_path: Path,
    timeout: int = 300,
    verbose: bool = False
) -> Tuple[bool, str]:
    """
    使用 IDA Pro decompile 一个二进制文件
    直接将脚本内容作为参数 passing
    """
    bin_name = binary_path.stem
    output_file = output_dir / f"{bin_name}.c"

    # 创建 log 目录
    log_dir = output_dir / "logs"
    log_dir.mkdir(parents=True, exist_ok=True)

    # 脚本内容 (内联)
    script_lines = [
        "import idc",
        "import idaapi",
        "import os",
        "",
        "bin_name = os.path.splitext(idc.get_input_file_path())[0]",
        f"output_dir = r\"{output_dir}\"",
        "output_file = os.path.join(output_dir, bin_name + \".c\")",
        "",
        "os.makedirs(output_dir, exist_ok=True)",
        "",
        "# Export C code",
        "result = idc.generate_file(idc.GFL_CDECOMP, output_file, 0, idc.BADADDR, 0)",
        "if result == 0:",
        "    print(f\"SUCCESS: {output_file}\")",
        "else:",
        "    print(f\"FAILED: {result}\")",
        "",
        "idaapi.qexit(0)"
    ]

    # 使用 -c 参数自动处理并退出
    # 使用 -A 静默模式
    # 使用 -S 执行脚本

    # 将脚本写入临时文件
    import tempfile
    script_fd, script_path = tempfile.mkstemp(suffix='.py', prefix='ida_export_')
    os.close(script_fd)
    Path(script_path).write_text('\\n'.join(script_lines))

    try:
        cmd = [
            str(idat_path),
            "-A",
            "-B",       # Batch mode, no UI
            f"-S{script_path}",
            str(binary_path)
        ]

        if verbose:
            print(f"  Processing: {bin_name}")

        result = subprocess.run(
            cmd,
            timeout=timeout,
            capture_output=True,
            text=True
        )

        if output_file.exists():
            return True, f"Generated {output_file}"
        else:
            stderr = result.stderr.strip() if result.stderr else result.stdout.strip()
            return False, f"Output file not created. {stderr}"

    except subprocess.TimeoutExpired:
        return False, f"Timeout after {timeout}s"
    except Exception as e:
        return False, f"Exception: {e}"
    finally:
        try:
            os.unlink(script_path)
        except:
            pass


# 完整版 - step by step
def decompile_binary_full(
    binary_path: Path,
    output_dir: Path,
    idat_path: Path,
    timeout: int = 300,
    verbose: bool = False
) -> Tuple[bool, str]:
    """
    完整的反编译流程:
    1. 使用 idat -A 批量处理
    2. IDAPython 脚本负责导出 C 代码
    """
    bin_name = binary_path.stem
    output_file = output_dir / f"{bin_name}.c"

    # 使用外部脚本文件
    script_path = Path(__file__).parent / "idapython_export_c.py"
    if not script_path.exists():
        return False, f"Script not found: {script_path}"

    # 设置环境变量传递输出目录
    env = os.environ.copy()
    env['IDA_OUTPUT_DIR'] = str(output_dir)

    try:
        # 使用 idat 批量模式
        cmd = [
            str(idat_path),
            "-A",       # Automatic mode
            f"-S{script_path}",
            str(binary_path)
        ]

        if verbose:
            print(f"  Processing: {bin_name}")

        result = subprocess.run(
            cmd,
            timeout=timeout,
            capture_output=True,
            text=True,
            env=env
        )

        if output_file.exists():
            return True, f"Generated {output_file}"
        else:
            stderr = result.stderr.strip() if result.stderr else result.stdout.strip()
            return False, f"Output not created: {stderr}"

    except subprocess.TimeoutExpired:
        return False, f"Timeout after {timeout}s"
    except Exception as e:
        return False, f"Exception: {e}"


def process_single_binary(args):
    """处理单个二进制文件的包装函数（用于 map）"""
    binary_path, output_dir, idat_path, timeout, verbose = args
    return decompile_binary_full(
        binary_path, output_dir, idat_path, timeout, verbose
    )


def main():
    parser = argparse.ArgumentParser(
        description="IDA Pro 批量反编译工具 - 使用 idat 导出二进制文件的 C 代码",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s -i ./binaries -o ./output
  %(prog)s -i /path/to/binary -o ./decompiled -p /Applications/IDA\\ Professional\\ 9.1.app/Contents/MacOS/idat
  %(prog)s -i ./binaries -o ./output -j 8 -t 600
        """
    )

    parser.add_argument('-i', '--input-dir', required=True,
                        help='要扫描的输入目录或二进制文件')
    parser.add_argument('-o', '--output-dir', required=True,
                        help='C 文件输出目录')
    parser.add_argument('-p', '--idat-path',
                        default='/Applications/IDA Professional 9.1.app/Contents/MacOS/idat',
                        help='idat 可执行文件路径')
    parser.add_argument('-j', '--jobs', type=int, default=4,
                        help='并行处理的进程数 (默认: 4)')
    parser.add_argument('-t', '--timeout', type=int, default=300,
                        help='每个文件的处理超时时间(秒) (默认: 300)')
    parser.add_argument('-v', '--verbose', action='store_true',
                        help='显示详细输出')

    args = parser.parse_args()

    # 验证 idat 路径
    idat_path = Path(args.idat_path)
    if not idat_path.exists():
        print(f"ERROR: idat not found at {idat_path}")
        sys.exit(1)

    # 验证输入
    input_path = Path(args.input_dir)
    if not input_path.exists():
        print(f"ERROR: Input path not found: {input_path}")
        sys.exit(1)

    # 创建输出目录
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # 查找二进制文件
    print(f"\n{'='*60}")
    print("IDA Pro 批量反编译工具")
    print(f"{'='*60}")
    print(f"\n输入目录: {input_path}")
    print(f"输出目录: {output_dir}")
    print(f"IDAT 路径: {idat_path}")
    print(f"并行进程: {args.jobs}")
    print(f"超时时间: {args.timeout}s")
    print()

    # 收集文件
    if input_path.is_file():
        binaries = [input_path]
        print(f"处理单个文件: {input_path}")
    else:
        print(f"扫描目录: {input_path}")
        binaries = find_binaries(input_path)

    if not binaries:
        print("\n未找到任何二进制文件")
        sys.exit(0)

    print(f"\n找到 {len(binaries)} 个二进制文件")

    # 处理文件
    print(f"\n{'='*60}")
    print("开始反编译...")
    print(f"{'='*60}\n")

    success_count = 0
    fail_count = 0

    if args.jobs > 1 and len(binaries) > 1:
        # 多进程并行处理
        import multiprocessing
        with multiprocessing.Pool(args.jobs) as pool:
            results = pool.map(process_single_binary, [
                (b, output_dir, idat_path, args.timeout, args.verbose)
                for b in binaries
            ])

            for binary, (success, msg) in zip(binaries, results):
                status = "OK" if success else "FAIL"
                print(f"[{status}] {binary.name}: {msg}")
                if success:
                    success_count += 1
                else:
                    fail_count += 1
    else:
        # 单进程处理
        for binary in binaries:
            success, msg = decompile_binary_full(
                binary, output_dir, idat_path, args.timeout, args.verbose
            )
            status = "OK" if success else "FAIL"
            print(f"[{status}] {binary.name}: {msg}")
            if success:
                success_count += 1
            else:
                fail_count += 1

    # 汇总
    print(f"\n{'='*60}")
    print("处理完成!")
    print(f"{'='*60}")
    print(f"成功: {success_count}")
    print(f"失败: {fail_count}")
    print(f"总计: {len(binaries)}")
    print(f"\n输出目录: {output_dir}")


if __name__ == "__main__":
    main()
