import os
import r2pipe
import re
import hashlib
import log_utils
import json


def calc_sha256(file_path):
    with open(file_path, 'rb') as f:
        bytes = f.read()
        sha256obj = hashlib.sha256(bytes)
        sha256 = sha256obj.hexdigest()
        return sha256

def extract_opcode(disasm_text):
    """
    从反汇编文本中提取操作码和操作数
    正则表达式用于匹配操作码和操作数，考虑到操作数可能包含空格和逗号
    """
    match = re.search(r"^\s*(\S+)(?:\s+(.*))?$", disasm_text)
    if match:
        opcode = match.group(1)
        # operands_str = match.group(2) if match.group(2) is not None else ""
        # split_pattern = re.compile(r",(?![^\[]*\])")  # 用于切分操作数的正则表达式
        # operands = split_pattern.split(operands_str)
        # return opcode, [op.strip() for op in operands if op.strip()]
        return opcode
    return ""

def get_graph_cfg_r2pipe(r2pipe_open):
    # CFG提取
    acfg_item = []
    try:
        # 获取函数列表
        function_list = r2pipe_open.cmdj("aflj")


        for function in function_list:
            # 局部函数内的特征提取
            node_list = []
            edge_list = []
            temp_edge_list = []
            block_list = r2pipe_open.cmdj("afbj @" + str(function['offset']))
            block_number = len(block_list)
            block_feature_list = []
            for block in block_list:
                node_list.append(block["addr"])

                # 获取基本块的反汇编指令
                disasm = r2pipe_open.cmdj("pdj " + str(block["ninstr"]) + " @" + str(block["addr"]))
                if disasm:
                    for op in disasm:
                        if op["type"] == "invalid":
                            continue
                        # TODO ：这里需要处理指令的特征提取
                        block_feature = ''
                        block_feature_list.append(block_feature)


                        # 处理跳转指令
                        if "jump" in op and op["jump"] != 0:
                            temp_edge_list.append([block["addr"], op["jump"]])
            for temp_edge in temp_edge_list:
                if temp_edge[1] in node_list:
                    edge_list.append(temp_edge)
            acfg = {
                'block_number': block_number,
                'block_edges': [[d[0] for d in edge_list], [d[1] for d in edge_list]],
                'block_features': block_feature_list
            }
            acfg_item.append(acfg)
        return True, "二进制可执行文件解析成功", acfg_item
    except Exception as e:
        return False, e, None

            # for block in block_list:
            #     node_list.append(block["addr"])
            #
            #     # 获取基本块的反汇编指令
            #     disasm = r2pipe_open.cmdj("pdj " + str(block["ninstr"]) + " @" + str(block["addr"]))
            #     node_info = []
            #     if disasm:
            #         for op in disasm:
            #             if op["type"] == "invalid":
            #                 continue
            #             opcode, operands = extract_opcode(op["disasm"])
            #             # 处理跳转指令
            #             if "jump" in op and op["jump"] != 0:
            #                 temp_edge_list.append([block["addr"], op["jump"]])
            #                 node_info.append([op["offset"], op["bytes"], opcode, op["jump"]])
            #             else:
            #                 node_info.append([op["offset"], op["bytes"], opcode, None])
            #     node_info_list.append(node_info)

        # 完成 CFG 构建后, 检查并清理不存在的出边


        # 获取排序后元素的原始索引
    #     sorted_indices = [i for i, v in sorted(enumerate(node_list), key=lambda x: x[1])]
    #     # 根据这些索引重新排列
    #     node_list = [node_list[i] for i in sorted_indices]
    #     node_info_list = [node_info_list[i] for i in sorted_indices]
    #
    #     return True, "二进制可执行文件解析成功", node_list, edge_list, node_info_list
    # except Exception as e:
    #     return False, e, None, None, None
def get_graph_fcg_r2pipe(r2pipe_open):
    # FCG提取
    try:
        function_list = r2pipe_open.cmdj("aflj")
        node_list = []
        func_name_list = []
        edge_list = []
        temp_edge_list = []
        function_num = len(function_list)

        for function in function_list:
            func_name_list.append(function["name"])
            r2pipe_open.cmd(f's ' + str(function["offset"]))
            pdf = r2pipe_open.cmdj('pdfj')
            if pdf is None:
                continue

            node_bytes = ""
            node_opcode = ""
            for op in pdf["ops"]:
                if op["type"] == "invalid":
                    continue

                node_bytes += op["bytes"]
                opcode = extract_opcode(op["disasm"])
                node_opcode += opcode + " "
                if "jump" in op and op["jump"] != 0:
                    temp_edge_list.append([function["offset"], op["jump"]])

            node_list.append(function["offset"])


        # 完成 FCG 构建后, 检查并清理不存在的出边
        for temp_edge in temp_edge_list:
            if temp_edge[1] in node_list:
                edge_list.append(temp_edge)
        sub_function_name_list = ('fcn.', 'loc.', 'main', 'entry')
        func_name_list = [func_name for func_name in func_name_list if not func_name.startswith(sub_function_name_list)]
        return True, "二进制可执行文件解析成功", function_num, edge_list, func_name_list
    except Exception as e:
        return False, e, None, None, None

def get_r2pipe(file_path):
    try:
        r2 = r2pipe.open(file_path, flags=['-2'])
        r2.cmd("aaa")
        r2.cmd('e arch=x86')
        return r2
    except Exception as e:
        return None

def init_logging():
    log_file = "./out/exe2json.log"
    logging = log_utils.setup_logger('exe2json', log_file)
    return logging


def exe_to_json(file_path, output_path):
    logging = init_logging()
    r2 = get_r2pipe(file_path)
    fcg_Operation_flag, fcg_Operation_message, function_num, function_fcg_edge_list, function_names = get_graph_fcg_r2pipe(r2)
    cfg_Operation_flag, cfg_Operation_message, cfg_item = get_graph_cfg_r2pipe(r2)
    file_fingerprint = calc_sha256(file_path)
    if fcg_Operation_flag and cfg_Operation_flag:
        json_obj = {
            'hash': file_fingerprint,
            'function_number': function_num,
            'function_edges': [[int(d[0]) for d in function_fcg_edge_list],
                               [int(d[1]) for d in function_fcg_edge_list]],
            'acfg_list': cfg_item,
            'function_names': function_names
        }
    else:
        logging.error(f"二进制可执行文件解析失败 文件地址{file_path}")
        if not fcg_Operation_flag:
            logging.error(f"fcg错误：{fcg_Operation_message}")
        if not cfg_Operation_flag:
            logging.error(f"cfg错误：{cfg_Operation_message}")
        return False
    r2.quit()
    result = json.dumps(json_obj,ensure_ascii=False)
    with open(os.path.join(output_path, file_fingerprint + '.jsonl'), 'w') as out:
        out.write(result)
        out.close()
    return True

if __name__ == '__main__':
    test_file_path = '/mnt/d/bishe/exe2json/sample/VirusShare_0a3b625380161cf92c4bb10135326bb5'
    exe_to_json(test_file_path, './out/json')