import subprocess
import os, time

# Для работы с IDA
from exetoi64 import make_i64
#

# Для работы с excel
from openpyxl import Workbook
from openpyxl.styles import Font, PatternFill
#

def main():

    # Задается вручную
    ida_path = r"C:\Users\user\Desktop\IDA Pro 7.7.220118 (Windows) (x86,x64,ARM64)\ida64.exe"
    bindiff_path = r"C:\tools\BinDiff\bin\bindiff.exe"
    exe_dir = r"C:\tools\temp"
    #

    # Создание .i64 из .exe файлов
    exe_list = find_exe_in_dir(exe_dir)
    for exe_path in exe_list:
        make_i64(ida_path, exe_path, exe_path + '.i64')
    #

    # Создание файлов экспорта, получение файлов сравнений
    start = time.perf_counter()
    create_exports(bindiff_path, exe_dir, exe_dir)
    compare_exports(bindiff_path, exe_dir)
    #

    # Получение имени файла с результатом сравнения
    exe_1_name = os.path.basename(exe_list[0])
    exe_2_name = os.path.basename(exe_list[1])
    result_file_name = fr'{exe_dir}\{exe_1_name}_vs_{exe_2_name}.results'
    #

    # Формирую список с инфой о сравнении функций
    compare_functions_info = parse_bindiff_log(result_file_name)
    #

    # Подсчет времени 
    elapsed = time.perf_counter() - start
    print(f"Elapsed: {elapsed}") 
    #

    # Выгрузка
    output_path = "output.xlsx"
    create_excel_from_bindiff_log(compare_functions_info, 
                                  output_path, 
                                  exe_1_name, 
                                  exe_2_name, 
                                  elapsed)
    # 

def create_excel_from_bindiff_log(compare_functions_info, 
                                  output_file, 
                                  exe_1_name, 
                                  exe_2_name, 
                                  elapsed):
    wb = Workbook()
    ws = wb.active
    ws.title = "BinDiff Results"

    # Заголовки столбцов
    headers = [
        f'Address in {exe_1_name}', 
        f'Address in {exe_2_name}', 
        'Match Type', 
        'Similarity', 
        'Confidence',
        'Basic Block Similarity', 
        'Flags', 
        'Additional Flags', 
        'Matching Algorithm',
        f'{exe_1_name}', 
        f'{exe_2_name}', 
        f'Matches: {len(compare_functions_info)}', 
        f'Elapsed: {elapsed:.2f} sec', 
        ' '
    ]
    ws.append(headers)

    bold_font = Font(bold=True)  # Стиль с жирным шрифтом
    for cell in ws[1]:  # ws[1] обращается к первой строке (заголовкам)
        cell.font = bold_font # Применяет стиль к каждой ячейке

    # Окраска ячеек с общей информацией
    green_fill = PatternFill(start_color="90EE00", end_color="90EE00", fill_type="solid")
    light_green_fill = PatternFill(start_color="90EE90", end_color="90EE90", fill_type="solid")
    ws['L1'].fill = green_fill
    ws['M1'].fill = light_green_fill

    # Добавляю данные
    for result in compare_functions_info:
        row = [
            result['addr1'], 
            result['addr2'], 
            result['match_type'],
            result['similarity'], 
            result['confidence'], 
            result['bb_similarity'],
            result['flags'], 
            result['additional_flags'], 
            result['matching_algorithm'],
            result['name1'], 
            result['name2'], 
            ' '
        ]
        ws.append(row)

    wb.save(output_file)

def parse_bindiff_log(log_file):
    results = []
    with open(log_file, 'r') as f:
        for line in f:
            fields = line.strip().split('\t')
            # в логе который создает bindiff по каждой функции 11 полей
            if len(fields) == 11:
                result = {
                    'addr1': fields[0], # адрес функции в первом бинаре
                    'addr2': fields[1], # адрес во втором 
                    'match_type': fields[2], # тип совпадения
                    'similarity': fields[3], # сходство 
                    'confidence': fields[4], # уверенность (на сколько bindiff уверен в правильности сопоставления)
                    'bb_similarity': fields[5], # сходство базовых блоков (показывает, насколько похожи базовые блоки внутри функции)
                    'flags': fields[6], # дополнительная информация
                    'additional_flags': fields[7], # дополнительная информация
                    'matching_algorithm': fields[8].replace("function: ", ""), # описание метода, который был использован для сопоставления функций
                    'name1': fields[9].strip('"'), # имя функции в первом бинаре
                    'name2': fields[10].strip('"') # имя во втором бинаре
                }
                results.append(result)
    return results

def find_exe_in_dir(directory):
    exe_files = []
    for root, _, files in os.walk(directory):
        for file in files:
            if file.endswith(".exe"):
                exe_files.append(os.path.join(root, file))
    return exe_files

def create_exports(bindiff_path, idb_dir, bindiff_exports_dir):
    cmd = f"{bindiff_path} --export {idb_dir} --output_dir {bindiff_exports_dir}"
    print('Creating export files...', end=' ')
    try:
        subprocess.run(cmd, capture_output=True, check=True)
        print('\033[92mSuccess!\033[0m')
    except subprocess.CalledProcessError as e:
        print(f"\033[91mError during export: {e}\033[0m")

def compare_exports(bindiff_path, bindiff_exports_dir):
    cmd = f"{bindiff_path} {bindiff_exports_dir} --output_format log"
    print('Comparing...', end=' ')
    try:
        subprocess.run(cmd, capture_output=True, check=True)
        print('\033[92mSuccess!\033[0m')
    except subprocess.CalledProcessError as e:
        print(f"\033[91mError during comparing: {e}\033[0m")

if __name__ == "__main__":
    main()