# ВНИМАНИЕ: работает с двумя бинарниками, !!!!! собранными под x64 !!!!!

import sqlite3
import os, time
import subprocess

# для работы с ida
from i64tosqlite import create_sqlite_from_i64
from exetoi64 import make_i64
#

# для работы с excel
import pandas as pd
from openpyxl.styles import PatternFill
from openpyxl import load_workbook
#

def main():

    # эти пути необходимо задавать вручную
    ida_executable = r"C:\tools\IDA\ida.exe"
    diaphora_dir = r"C:\tools\diaphora"
    exe_1_path = r"C:\tools\my.exe"
    exe_2_path = r"C:\tools\ilya.exe"
    #

    # генерация i64 из exe
    start = time.perf_counter()
    i64_1_path = exe_1_path + ".i64"
    i64_2_path = exe_2_path + ".i64"
    make_i64(ida_executable, exe_1_path, i64_1_path)
    make_i64(ida_executable, exe_2_path, i64_2_path)
    #

    # генерация sqlite базы из i64
    diaphora_script_path = diaphora_dir + r"\diaphora_ida.py"
    diaphora_path = diaphora_dir + r"\diaphora.py"
    sqlite1_path = i64_1_path+'.sqlite'
    sqlite2_path = i64_2_path+'.sqlite'
    create_sqlite_from_i64(ida_executable, i64_1_path, diaphora_script_path, sqlite1_path)
    create_sqlite_from_i64(ida_executable, i64_2_path, diaphora_script_path, sqlite2_path)
    #

    # создание выходного файла
    sql_output_path = r'result.sqlite'
    excel_output_path = r'output.xlsx'
    generate_output_db(diaphora_path, sqlite1_path, sqlite2_path, sql_output_path)
    result_data = read_results_table(sql_output_path)
    elapsed = time.perf_counter() - start
    #

    # удаляю ненужное
    os.remove(i64_1_path)
    os.remove(i64_2_path)
    os.remove(sqlite1_path)
    os.remove(sqlite2_path)
    os.remove(sql_output_path)
    #

    # получаем имена файлов без пути
    exe_1_name = os.path.basename(exe_1_path)
    exe_2_name = os.path.basename(exe_2_path)
    #

    export_to_excel(result_data, excel_output_path, elapsed, exe_1_name, exe_2_name)
    print(f"\033[92mResults exported to\033[0m {excel_output_path}")
    print(f'\033[92mAnalysis time\033[0m {elapsed} sec')
    os.startfile(excel_output_path)

def generate_output_db(diaphora_path, sqlite1_path, sqlite2_path, output_path):
    cmd = f'py {diaphora_path} {sqlite1_path} {sqlite2_path} -o {output_path}'
    try:
        # capture_output=True подавляет вывод сообщений от субпроцесса
        subprocess.run(cmd, check=True, capture_output=True)
    except subprocess.CalledProcessError as e:
        print(f"\033[31mError subprocess run: {e.stderr}\033[0m")

def read_results_table(output_path):
    conn = sqlite3.connect(output_path)
    cursor = conn.cursor()
    cursor.execute("SELECT * FROM results")
    rows = cursor.fetchall()
    columns = [description[0] for description in cursor.description]
    result_list = []
    for row in rows:
        result_list.append(dict(zip(columns, row)))
    
    conn.close()
    return result_list

def export_to_excel(result_data, excel_output_path, elapsed_time, exe_1_name, exe_2_name):
    df = pd.DataFrame(result_data)
    df = df[['type', 'name', 'name2', 'description']]
    
    # Переименование столбцов
    df.rename(columns={'name': exe_1_name, 'name2': exe_2_name}, inplace=True)
    
    df['matches'] = ' '
    df['elapsed'] = ' '
    df.at[0, 'matches'] = len(result_data)
    df.at[0, 'elapsed'] = f'{elapsed_time:.2f} sec'
    df.to_excel(excel_output_path, index=False)
    
    # Форматирование Excel
    wb = load_workbook(excel_output_path)
    ws = wb.active
    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['E2'].fill = green_fill
    ws['F2'].fill = light_green_fill
    wb.save(excel_output_path)

if __name__ == "__main__":
    main()