DBs/Dataset-1-CodeCMR/Dataset-1-CodeCMR creation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64143f99",
   "metadata": {},
   "outputs": [],
   "source": [
    "##############################################################################\n",
    "#                                                                            #\n",
    "#  Code for the USENIX Security '22 paper:                                   #\n",
    "#  How Machine Learning Is Solving the Binary Function Similarity Problem.   #\n",
    "#                                                                            #\n",
    "#  MIT License                                                               #\n",
    "#                                                                            #\n",
    "#  Copyright (c) 2019-2022 Cisco Talos                                       #\n",
    "#                                                                            #\n",
    "#  Permission is hereby granted, free of charge, to any person obtaining     #\n",
    "#  a copy of this software and associated documentation files (the           #\n",
    "#  \"Software\"), to deal in the Software without restriction, including       #\n",
    "#  without limitation the rights to use, copy, modify, merge, publish,       #\n",
    "#  distribute, sublicense, and/or sell copies of the Software, and to        #\n",
    "#  permit persons to whom the Software is furnished to do so, subject to     #\n",
    "#  the following conditions:                                                 #\n",
    "#                                                                            #\n",
    "#  The above copyright notice and this permission notice shall be            #\n",
    "#  included in all copies or substantial portions of the Software.           #\n",
    "#                                                                            #\n",
    "#  THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND,           #\n",
    "#  EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF        #\n",
    "#  MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND                     #\n",
    "#  NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE    #\n",
    "#  LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION    #\n",
    "#  OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION     #\n",
    "#  WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.           #\n",
    "#                                                                            #\n",
    "#  Dataset-1-CodeCMR creation                                                #\n",
    "#                                                                            #\n",
    "##############################################################################"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "594087b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "import os\n",
    "import pandas as pd\n",
    "import random\n",
    "import shutil\n",
    "\n",
    "from collections import Counter\n",
    "from collections import defaultdict\n",
    "from itertools import chain\n",
    "from itertools import compress\n",
    "from tqdm import tqdm\n",
    "\n",
    "pdcsv = lambda x: pd.read_csv(x, index_col=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46e7625c",
   "metadata": {},
   "source": [
    "The following table summarizes the criteria used to generate positive pairs for each task:\n",
    "* The `X` indicates that the variable is required to be different in each pair\n",
    "* The `*` indicates that the variable is free and may differ (but it isn't required).\n",
    "\n",
    "```\n",
    "|       | Architecture | Bitness | Compiler | Version | Optimization |\n",
    "|-------|--------------|---------|----------|---------|--------------|\n",
    "| arch  | X            |         |          |         |              |\n",
    "| bit   |              | X       |          |         |              |\n",
    "| comp  |              |         | X        | X       |              |\n",
    "| ver   |              |         |          | X       |              |\n",
    "| opt   |              |         |          |         | X            |\n",
    "| XA    | X            | X       |          |         |              |\n",
    "| XA+XO | X            | X       |          |         | X            |\n",
    "| XC    |              |         | X        | X       | X            |\n",
    "| XC+XB |              | X       | X        | X       | X            |\n",
    "| XM    | *            | *       | *        | *       | *            |\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2b41299",
   "metadata": {},
   "outputs": [],
   "source": [
    "CATEGORIES = [\n",
    "    \"project\",\n",
    "    \"library\",\n",
    "    \"arch\",\n",
    "    \"bit\",\n",
    "    \"compiler\",\n",
    "    \"version\",\n",
    "    \"optimizations\",\n",
    "]\n",
    "\n",
    "TASKS_DICT = {\n",
    "    # For any positive pair, the project and the library are the same.\n",
    "    #   True: the variable is required to have the same value in the positive pair\n",
    "    #   False: the variable is required to have different values in the negative pair.\n",
    "    \"arch\": [\n",
    "        True, True, False, True, True, True, True],\n",
    "    \"bit\": [\n",
    "        True, True, True, False, True, True, True],\n",
    "    \"comp\": [\n",
    "        True, True, True, True, False, False, True],\n",
    "    \"ver\": [\n",
    "        True, True, True, True, True, False, True],\n",
    "    \"opt\": [\n",
    "        True, True, True, True, True, True, False],\n",
    "    \"XA\": [\n",
    "        True, True, False, False, True, True, True],\n",
    "    \"XA+XO\": [\n",
    "        True, True, False, False, True, True, False],\n",
    "    \"XC\": [\n",
    "        True, True, True, True, False, False, False],\n",
    "    \"XC+XB\": [\n",
    "        True, True, True, False, False, False, False],\n",
    "    # The following would be the XA+XC test\n",
    "    # \"XA+XC\": [\n",
    "    #    True, True, False, False, False, False, False]\n",
    "}\n",
    "\n",
    "# The XO test is the same as the opt one.\n",
    "TASKS_DICT[\"XO\"] = TASKS_DICT[\"opt\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72437f9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "DATASET_ONE_DICT = {\n",
    "    \"projects\": {\n",
    "        \"training\": [\"openssl\", \"clamav\", \"curl\", \"unrar\"],\n",
    "        \"validation\": [\"zlib\"],\n",
    "        \"test\": [\"z3\", \"nmap\"],\n",
    "    },\n",
    "    \"eval\": {\n",
    "        \"validation\": {\n",
    "            \"similarity\": {\"XA\": 10000, \"XC\": 10000, \"XC+XB\": 10000, \"XM\": 10000}\n",
    "        },\n",
    "        \"test\": {\n",
    "            \"similarity\": {\n",
    "                \"XA\": 50000,\n",
    "                \"XC\": 50000,\n",
    "                \"XC+XB\": 50000,\n",
    "                \"XM\": 50000,\n",
    "            },\n",
    "            \"rank\": {\"XA\": 200, \"XC\": 200, \"XC+XB\": 200, \"XM\": 200},\n",
    "        },\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c59899a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# where to save the new dataset\n",
    "OUTPUT_DIR = \"../Dataset-1-CodeCMR-new/\"\n",
    "\n",
    "if not os.path.isdir(OUTPUT_DIR):\n",
    "    os.makedirs(OUTPUT_DIR)\n",
    "    print(f\"[D] DIR created: {OUTPUT_DIR}\")\n",
    "    \n",
    "for dirname in ['validation', 'testing']:\n",
    "    tmp_path = os.path.join(OUTPUT_DIR, \"pairs\", dirname)\n",
    "    if not os.path.isdir(tmp_path):\n",
    "        os.makedirs(tmp_path)\n",
    "        print(f\"[D] DIR created: {tmp_path}\")\n",
    "\n",
    "for dirname in ['training', 'validation', 'testing']:\n",
    "    tmp_path = os.path.join(OUTPUT_DIR, \"features\", dirname)\n",
    "    if not os.path.isdir(tmp_path):\n",
    "        os.makedirs(tmp_path)\n",
    "        print(f\"[D] DIR created: {tmp_path}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7bf4492f",
   "metadata": {},
   "source": [
    "### Create a training / validation / test split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d6f96b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The starting point\n",
    "CSV_FLOWCHART_FP = \"../Dataset-1/features/flowchart_Dataset-1.csv\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fa99c91",
   "metadata": {},
   "source": [
    "Summary:\n",
    "   * Step 0 - Read the list of functions from the output of IDA flowchart\n",
    "   * Step 1 -  Filter the functions with less than 5 BBs\n",
    "   * Step 2 - Remove duplicated hashopcodes to remove duplicated functions\n",
    "   * Step 3 - Extract compilation variables from idb_path\n",
    "   * Step 4 - Create training, validation and test splits\n",
    "   * Step 5 - Remove common function names across splits\n",
    "   * Step 6 - Remove singleton functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1e1b7ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step0 - Read the list of functions from the output of IDA flowchart\n",
    "df = pd.read_csv(CSV_FLOWCHART_FP)\n",
    "print(f\"Shape: {df.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85dda515",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Remove the column with the list of basic-blocks\n",
    "del df['bb_list']\n",
    "print(f\"Shape: {df.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e7110bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step1 -  Filter the functions with less than 5 BBs\n",
    "df = df[df['bb_num'] >= 5]\n",
    "print(f\"Shape: {df.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef52c63c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step2 - Remove duplicated hashopcodes to remove duplicated functions\n",
    "df.drop_duplicates('hashopcodes', keep='first', inplace=True)\n",
    "print(f\"Shape: {df.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b78d5ac5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step3 - Extract compilation variables from idb_path\n",
    "compilation_var = list()\n",
    "for path in df['idb_path']:\n",
    "    slist = path.split(\"/\")[2:]\n",
    "    project = slist[0]\n",
    "    slist = slist[1].split(\"_\")\n",
    "    library = slist[1].replace(\".i64\", \"\")\n",
    "    arch, comp, ver, opt = slist[0].split(\"-\")\n",
    "    bit = \"32\" if \"32\" in arch.replace(\"86\", \"32\") else \"64\"\n",
    "    arch = arch.replace(\"32\", \"\").replace(\"64\", \"\").replace(\"86\", \"\")\n",
    "    if comp == \"gcc\":\n",
    "        ver = \"gcc_\" + ver\n",
    "    pkl =  \"_\".join(path.split(\"/\")[2:]).replace(\".i64\", \"_codeCMR.pkl\")\n",
    "    compilation_var.append([project, library, arch, bit, comp, ver, opt, pkl])\n",
    "\n",
    "    \n",
    "# Convert to NumPy Array\n",
    "compilation_var = np.array(compilation_var)\n",
    "\n",
    "# Add compilation variables to the DataFrame\n",
    "df['project'] = compilation_var[:,0].tolist()\n",
    "df['library'] = compilation_var[:,1].tolist()\n",
    "df['arch'] = compilation_var[:,2].tolist()\n",
    "df['bit'] = compilation_var[:,3].tolist()\n",
    "df['compiler'] = compilation_var[:,4].tolist()\n",
    "df['version'] = compilation_var[:,5].tolist()\n",
    "df['optimizations'] = compilation_var[:,6].tolist()\n",
    "df['pickle_path'] = compilation_var[:,7].tolist()\n",
    "\n",
    "print(f\"Shape: {df.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58ceeb39",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optional - Remove unsupported architectures\n",
    "# We didn't have access to the Hex-Rays decompiler for MIPS\n",
    "df = df[df['arch'] != 'mips']\n",
    "df.reset_index(inplace=True, drop=True)\n",
    "print(f\"Shape: {df.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec3553a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step4 - Create training, validation and test splits\n",
    "df_training = df[df['project'].isin(DATASET_ONE_DICT['projects']['training'])]\n",
    "df_training['pickle_path'] = df_training['pickle_path'].apply(lambda x: \"features/training/\" + x)\n",
    "print(f\"Shape df_training: \\t{df_training.shape}\")\n",
    "\n",
    "df_validation = df[df['project'].isin(DATASET_ONE_DICT['projects']['validation'])]\n",
    "df_validation['pickle_path'] = df_validation['pickle_path'].apply(lambda x: \"features/validation/\" + x)\n",
    "print(f\"Shape df_validation: \\t{df_validation.shape}\")\n",
    "\n",
    "df_test = df[df['project'].isin(DATASET_ONE_DICT['projects']['test'])]\n",
    "df_test['pickle_path'] = df_test['pickle_path'].apply(lambda x: \"features/testing/\" + x)\n",
    "print(f\"Shape df_test: \\t\\t{df_test.shape}\")\n",
    "\n",
    "# Reset indexes\n",
    "df_training.reset_index(inplace=True, drop=True)\n",
    "df_validation.reset_index(inplace=True, drop=True)\n",
    "df_test.reset_index(inplace=True, drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ae8de8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step5 - Remove common function names across splits\n",
    "\n",
    "# Check for common function names in training and test\n",
    "r1 = set(df_training['func_name'].values) & set(df_test['func_name'].values)\n",
    "print(f\"# function names to remove: {len(r1)} (train & test)\")\n",
    "\n",
    "# Check for common function names in training and validation\n",
    "r2 = set(df_training['func_name'].values) & set(df_validation['func_name'].values)\n",
    "print(f\"# function names to remove: {len(r2)} (train & validation)\")\n",
    "\n",
    "df_training = df_training[~df_training['func_name'].isin(r1 | r2)]\n",
    "df_training.reset_index(inplace=True, drop=True)\n",
    "print(f\"Shape df_training: \\t{df_training.shape}\")\n",
    "\n",
    "print()\n",
    "\n",
    "# Check for common function names in validation and test\n",
    "r3 = set(df_validation['func_name'].values) & set(df_test['func_name'].values)\n",
    "print(f\"# function names to remove: {len(r3)} (validation & test)\")\n",
    "\n",
    "df_test = df_test[~df_test['func_name'].isin(r3)]\n",
    "df_test.reset_index(inplace=True, drop=True)\n",
    "print(f\"Shape df_test: \\t\\t{df_test.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef50eeb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step6 - Remove singleton functions\n",
    "for df_t in [df_training, df_validation, df_test]:\n",
    "    sl = [x for x, y in df_t[[\"library\", \"func_name\"]].value_counts().items() if y < 2]\n",
    "    gg = df_t.groupby([\"library\", \"func_name\"]).groups\n",
    "    idx_list = list(chain(*[list(gg[i]) for i in sl]))\n",
    "    print(f\"[D] # function to remove: {len(idx_list)}\")\n",
    "    \n",
    "    df_t.drop(idx_list, inplace=True)\n",
    "    df_t.reset_index(inplace=True, drop=True)\n",
    "    print(f\"[D] Shape: {df_t.shape}\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3f852f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Shape df_training: \\t{df_training.shape}\")\n",
    "print(f\"Shape df_validation: \\t{df_validation.shape}\")\n",
    "print(f\"Shape df_test: \\t\\t{df_test.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa96b47b",
   "metadata": {},
   "source": [
    "### Create positive and negative pairs for validation and test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ba7c0d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_similarity_pairs(df_input, num_pairs, test):\n",
    "    \"\"\"\n",
    "    Generate \"num_pairs\" positive function pairs by sub sampling all the\n",
    "    possible function combinations. Use this function when the number\n",
    "    of ((libraries, function_names)) is limited to few hundreds.\n",
    "    \"\"\"\n",
    "    # Map (libraries, function_names) to the indexes in the DB\n",
    "    libfunc_dict = {\n",
    "        k: list(v) for k, v in df_input.groupby([\"library\", \"func_name\"]).groups.items()\n",
    "    }\n",
    "\n",
    "    pos_pair_set = set()\n",
    "    neg_pair_set = set()\n",
    "    pos_pair_list = list()\n",
    "    neg_pair_list = list()\n",
    "\n",
    "    # Iterate over each library/func_name pair\n",
    "    for entry in tqdm(libfunc_dict.keys(), ncols=100):\n",
    "        libname, fname = entry\n",
    "\n",
    "        # Get the list of indexes associated to the ((libname, fname)) pair\n",
    "        idx_libfunc = libfunc_dict[entry]\n",
    "        # DataFrame for the library/func_name pair\n",
    "        df_libfunc = df_input.iloc[idx_libfunc]\n",
    "\n",
    "        # Get the list of indexes to select negative pairs\n",
    "        idx_list_neg = df_input[df_input[\"func_name\"] != fname].index\n",
    "\n",
    "        # (<-- left) Iterate over each function for the ((libname, fname)) pair\n",
    "        for idx_left_p in idx_libfunc:\n",
    "\n",
    "            # Extract the compilation variables\n",
    "            comp_data = df_input.iloc[idx_left_p][CATEGORIES].values\n",
    "\n",
    "            # For the XM test, any combination is valid\n",
    "            idx_list_pos = idx_libfunc\n",
    "\n",
    "            if test != \"XM\":\n",
    "                mask = TASKS_DICT[test]\n",
    "                # Build the constraints dict\n",
    "                #   if m is True: the variable is required to be the same in the positive pair\n",
    "                fd = {c: v for m, c, v in zip(mask, CATEGORIES, comp_data) if m}\n",
    "                constraints = [(df_libfunc[k] == v) for k, v in fd.items()]\n",
    "                #   if m is False: the variable is required to be different in the positive pair.\n",
    "                fd = {c: v for m, c, v in zip(mask, CATEGORIES, comp_data) if not m}\n",
    "                constraints += [(df_libfunc[k] != v) for k, v in fd.items()]\n",
    "\n",
    "                # Get the list of indexes of candidate right functions to generate positive pairs\n",
    "                idx_list_pos = df_libfunc[np.logical_and.reduce(constraints)].index\n",
    "\n",
    "            # Remove the left function from the list\n",
    "            idx_list_pos = [idx for idx in idx_list_pos if idx != idx_left_p]\n",
    "\n",
    "            # Iterate over each (--> right) function\n",
    "            for idx_right_p in idx_list_pos:\n",
    "                pos_pair = (idx_left_p, idx_right_p)\n",
    "\n",
    "                # Check if the pos_pair is already in the list\n",
    "                if tuple(sorted(pos_pair)) not in pos_pair_set:\n",
    "                    pos_pair_set.add(tuple(sorted(pos_pair)))\n",
    "                    pos_pair_list.append(pos_pair)\n",
    "\n",
    "                    # Generate the corresponding negative pair\n",
    "                    is_success = False\n",
    "                    while not is_success:\n",
    "                        idx_right_n = random.choice(idx_list_neg)\n",
    "                        neg_pair = (idx_left_p, idx_right_n)\n",
    "\n",
    "                        # Check if the neg_pair is already in the list\n",
    "                        if tuple(sorted(neg_pair)) not in neg_pair_set:\n",
    "                            neg_pair_set.add(tuple(sorted(neg_pair)))\n",
    "                            neg_pair_list.append(neg_pair)\n",
    "                            is_success = True\n",
    "\n",
    "    # print(\n",
    "    #     f\"[D] Before sampling - pos: {len(pos_pair_list)} - neg: {len(neg_pair_list)}\"\n",
    "    # )\n",
    "\n",
    "    # Sub sample the positive and negative pairs to num_pairs\n",
    "    if len(pos_pair_list) > num_pairs:\n",
    "        sampled_list = random.sample(list(range(len(pos_pair_list))), num_pairs)\n",
    "        pos_pair_list = [pos_pair_list[x] for x in sampled_list]\n",
    "        neg_pair_list = [neg_pair_list[x] for x in sampled_list]\n",
    "        # print(\n",
    "        #     f\"[D] After sampling - pos: {len(pos_pair_list)} - neg: {len(neg_pair_list)}\"\n",
    "        # )\n",
    "\n",
    "    return pos_pair_list, neg_pair_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "002e39e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_similarity_pairs_random_version(df_input, num_pairs, test, num_negatives=1):\n",
    "    \"\"\"\n",
    "    Randomly generate \"num_pairs\" positive function pairs. Use this function\n",
    "    when the number of (libraries, function_names) pairs is > 1 thousand.\n",
    "    \"\"\"\n",
    "    # Map (libraries, function_names) to the indexes in the DB\n",
    "    libfunc_dict = {\n",
    "        k: list(v) for k, v in df_input.groupby([\"library\", \"func_name\"]).groups.items()\n",
    "    }\n",
    "    libfunc_list = list(libfunc_dict.keys())\n",
    "\n",
    "    pos_pair_set = set()\n",
    "    neg_pair_set = set()\n",
    "    pos_pair_list = list()\n",
    "    neg_pair_list = list()\n",
    "\n",
    "    with tqdm(total=num_pairs, ncols=100) as pbar:\n",
    "        # Iterate num_pairs time to create the pos/neg function pairs\n",
    "        for _ in range(num_pairs):\n",
    "\n",
    "            # Iterate until a positive function pair is generated\n",
    "            is_success_pos = False\n",
    "            while not is_success_pos:\n",
    "\n",
    "                # Randomly select a library/func_name pair\n",
    "                entry = random.choice(libfunc_list)\n",
    "                libname, fname = entry\n",
    "                # Get the list of indexes associated to the library/func_name pair\n",
    "                idx_libfunc = libfunc_dict[entry]\n",
    "                # DataFrame for the library/func_name pair\n",
    "                df_libfunc = df_input.iloc[idx_libfunc]\n",
    "\n",
    "                # Randomly select a (<-- left) function\n",
    "                idx_left_p = random.choice(idx_libfunc)\n",
    "                # Extract the compilation variables\n",
    "                comp_data = df_input.iloc[idx_left_p][CATEGORIES].values\n",
    "\n",
    "                # For the XM test, any combination is valid\n",
    "                idx_list_pos = idx_libfunc\n",
    "\n",
    "                if test != \"XM\":\n",
    "                    mask = TASKS_DICT[test]\n",
    "                    # Build the constraints dict\n",
    "                    #   if m is True: the variable is required to be the same in the positive pair\n",
    "                    fd = {c: v for m, c, v in zip(mask, CATEGORIES, comp_data) if m}\n",
    "                    constraints = [(df_libfunc[k] == v) for k, v in fd.items()]\n",
    "                    #   if m is False: the variable is required to be different in the positive pair.\n",
    "                    fd = {c: v for m, c, v in zip(mask, CATEGORIES, comp_data) if not m}\n",
    "                    constraints += [(df_libfunc[k] != v) for k, v in fd.items()]\n",
    "\n",
    "                    # Get the list of indexes of candidate right functions to generate positive pairs\n",
    "                    idx_list_pos = df_libfunc[np.logical_and.reduce(constraints)].index\n",
    "\n",
    "                # Remove the left function from the list\n",
    "                idx_list_pos = [idx for idx in idx_list_pos if idx != idx_left_p]\n",
    "\n",
    "                # No functions are left. Retry\n",
    "                if len(idx_list_pos) == 0:\n",
    "                    continue\n",
    "\n",
    "                # Randomly select a (<-- right) function\n",
    "                idx_right_p = random.choice(idx_list_pos)\n",
    "                pos_pair = (idx_left_p, idx_right_p)\n",
    "                if tuple(sorted(pos_pair)) not in pos_pair_set:\n",
    "                    pos_pair_set.add(tuple(sorted(pos_pair)))\n",
    "                    pos_pair_list.append(pos_pair)\n",
    "                    is_success_pos = True\n",
    "\n",
    "                    for _ in range(num_negatives):\n",
    "                        # Generate the corresponding negative pair\n",
    "                        is_success_neg = False\n",
    "                        while not is_success_neg:\n",
    "                            idx_right_n = random.randint(0, df_input.shape[0] - 1)\n",
    "                            if df_input.iloc[idx_right_n][\"func_name\"] == fname:\n",
    "                                continue\n",
    "                            neg_pair = (idx_left_p, idx_right_n)\n",
    "\n",
    "                            # Check if the neg_pair is already in the list\n",
    "                            if tuple(sorted(neg_pair)) not in neg_pair_set:\n",
    "                                neg_pair_set.add(tuple(sorted(neg_pair)))\n",
    "                                neg_pair_list.append(neg_pair)\n",
    "                                is_success_neg = True\n",
    "\n",
    "            # Update the progress bar\n",
    "            pbar.update(1)\n",
    "\n",
    "        # print(f\"[D] pos: {len(pos_pair_list)} - neg: {len(neg_pair_list)}\")\n",
    "\n",
    "    return pos_pair_list, neg_pair_list\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e2adc71",
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_dicts_into_dataframes(df_input, dataset_dict):\n",
    "    pair_columns = [\n",
    "        \"idb_path_1\",\n",
    "        \"pickle_path_1\",\n",
    "        \"fva_1\",\n",
    "        \"func_name_1\",\n",
    "        \"idb_path_2\",\n",
    "        \"pickle_path_2\",\n",
    "        \"fva_2\",\n",
    "        \"func_name_2\",\n",
    "        \"db_type\",\n",
    "    ]\n",
    "\n",
    "    pos_pair_dict = defaultdict(list)\n",
    "    neg_pair_dict = defaultdict(list)\n",
    "    \n",
    "    # Iterate over each positive and negative pair.\n",
    "    #   Select the required info and save it in a new dictionary.\n",
    "    for task in dataset_dict:\n",
    "        for pos_pair in dataset_dict[task][\"pos\"]:\n",
    "            for c in [\"idb_path\", \"pickle_path\", \"fva\", \"func_name\"]:\n",
    "                pos_pair_dict[c + \"_1\"].append(df_input.iloc[pos_pair[0]][c])\n",
    "                pos_pair_dict[c + \"_2\"].append(df_input.iloc[pos_pair[1]][c])\n",
    "            pos_pair_dict[\"db_type\"].append(task)\n",
    "\n",
    "        for neg_pair in dataset_dict[task][\"neg\"]:\n",
    "            for c in [\"idb_path\", \"pickle_path\", \"fva\", \"func_name\"]:\n",
    "                neg_pair_dict[c + \"_1\"].append(df_input.iloc[neg_pair[0]][c])\n",
    "                neg_pair_dict[c + \"_2\"].append(df_input.iloc[neg_pair[1]][c])\n",
    "            neg_pair_dict[\"db_type\"].append(task)\n",
    "    \n",
    "    # Convert the local pair_dicts into DataFrames\n",
    "    df_pos = pd.DataFrame.from_dict(pos_pair_dict)\n",
    "    df_neg = pd.DataFrame.from_dict(neg_pair_dict)\n",
    "    \n",
    "    # Check/change the order of the columns\n",
    "    df_pos = df_pos[pair_columns]\n",
    "    df_neg = df_neg[pair_columns]\n",
    "    return df_pos, df_neg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1498aa5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_summary(dataset_dict):\n",
    "    print(\"[D] Summary:\") \n",
    "    for task in dataset_dict:\n",
    "        print(\n",
    "            \"[D] \\tTask: {:5} - pos: {:5} neg: {:5}\".format(\n",
    "                task, len(dataset_dict[task][\"pos\"]), len(dataset_dict[task][\"neg\"])\n",
    "            )\n",
    "        )\n",
    "    print(\"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60a6c6fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_free_variables(df_input, task_list, dataset_dict):\n",
    "    for task in task_list:\n",
    "        # Skip \"XM\"\n",
    "        if task not in TASKS_DICT:\n",
    "            continue\n",
    "\n",
    "        print(\"-\" * 100 + \"\\n\")\n",
    "        print(f\"[D] Task: {task}\\n\")\n",
    "\n",
    "        # Get the name of the free variables for each task\n",
    "        free_variables = list(\n",
    "            compress(CATEGORIES, [not x for x in TASKS_DICT[task]])\n",
    "        )\n",
    "\n",
    "        v_list = list()\n",
    "        for pos_pair in dataset_dict[task][\"pos\"]:\n",
    "            # Get the values associated to the free variables\n",
    "            vv = df_input.iloc[list(pos_pair)][free_variables].values\n",
    "            # Sort them to avoid counting the permutations\n",
    "            vv = tuple(sorted([tuple(x) for x in vv]))\n",
    "            v_list.append(vv)\n",
    "\n",
    "        # Print the frequency of each combination\n",
    "        for k, v in Counter(v_list).most_common():\n",
    "            print(f\"\\t{v:5}, {k}\")\n",
    "\n",
    "        print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f9ab051",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_pos_neg_dataset(\n",
    "    df_input, task_dict, output_dir, output_fs, rand=True, num_negatives=1\n",
    "):\n",
    "    print(\"[D] Creating the pos/neg function pairs...\", flush=True)\n",
    "    dataset_dict = defaultdict(dict)\n",
    "\n",
    "    for task, num_pairs in task_dict.items():\n",
    "        ppl, npl = None, None\n",
    "        if rand:\n",
    "            # Use the random version of the pair generation function\n",
    "            ppl, npl = create_similarity_pairs_random_version(\n",
    "                df_input, num_pairs, task, num_negatives\n",
    "            )\n",
    "        else:\n",
    "            ppl, npl = create_similarity_pairs(df_input, num_pairs, task)\n",
    "        dataset_dict[task][\"pos\"] = ppl\n",
    "        dataset_dict[task][\"neg\"] = npl\n",
    "\n",
    "    print_summary(dataset_dict)\n",
    "\n",
    "    print(\"[D] Converting the positive/negative pairs into CSV...\", flush=True)\n",
    "    df_pos, df_neg = convert_dicts_into_dataframes(df_input, dataset_dict)\n",
    "\n",
    "    pos_fp = os.path.join(output_dir, output_fs.format(\"pos\"))\n",
    "    df_pos.to_csv(pos_fp)\n",
    "    print(f\"[D] \\tPos CSV: {pos_fp}\")\n",
    "\n",
    "    neg_fp = os.path.join(output_dir, output_fs.format(\"neg\"))\n",
    "    df_neg.to_csv(neg_fp)\n",
    "    print(f\"[D] \\tNeg CSV: {neg_fp}\")\n",
    "\n",
    "    # For debug only\n",
    "    print_free_variables(df_input, task_dict.keys(), dataset_dict)\n",
    "    \n",
    "    selected_functions = set()\n",
    "    for task in dataset_dict:\n",
    "        for pair in dataset_dict[task][\"pos\"]:\n",
    "            selected_functions.update(list(pair))\n",
    "        for pair in dataset_dict[task][\"neg\"]:\n",
    "            selected_functions.update(list(pair))\n",
    "    return selected_functions\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0440be0",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Create pairs for validation dataset\n",
    "sf_set = create_pos_neg_dataset(\n",
    "    df_validation,\n",
    "    DATASET_ONE_DICT[\"eval\"][\"validation\"][\"similarity\"],\n",
    "    os.path.join(OUTPUT_DIR, \"pairs\", \"validation\"),\n",
    "    \"{}_validation_Dataset-1-CodeCMR.csv\",\n",
    "    rand=False,\n",
    "    num_negatives=1\n",
    ")\n",
    "\n",
    "df_validation = df_validation.iloc[list(sf_set)]\n",
    "df_validation.reset_index(inplace=True, drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ed95d38",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create pairs for test dataset\n",
    "sf_set_1 = create_pos_neg_dataset(\n",
    "    df_test,\n",
    "    DATASET_ONE_DICT[\"eval\"][\"test\"][\"similarity\"],\n",
    "    os.path.join(OUTPUT_DIR, \"pairs\", \"testing\"),\n",
    "    \"{}_testing_Dataset-1-CodeCMR.csv\",\n",
    "    rand=True,\n",
    "    num_negatives=1\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae9598dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create pairs for test rank dataset\n",
    "sf_set_2 = create_pos_neg_dataset(\n",
    "    df_test,\n",
    "    DATASET_ONE_DICT[\"eval\"][\"test\"][\"rank\"],\n",
    "    os.path.join(OUTPUT_DIR, \"pairs\", \"testing\"),\n",
    "    \"{}_rank_testing_Dataset-1-CodeCMR.csv\",\n",
    "    rand=True,\n",
    "    num_negatives=100\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "febe694c",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_test = df_test.iloc[list(sf_set_1 | sf_set_2)]\n",
    "df_test.reset_index(inplace=True, drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba69269a",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Shape df_training: \\t{df_training.shape}\")\n",
    "print(f\"Shape df_validation: \\t{df_validation.shape}\")\n",
    "print(f\"Shape df_test: \\t\\t{df_test.shape}\")\n",
    "\n",
    "# Save the \"selected functions\" to a CSV.\n",
    "# This will be useful to post-process the results.\n",
    "df_validation.to_csv(os.path.join(OUTPUT_DIR, \"validation_Dataset-1-CodeCMR.csv\"))\n",
    "df_training.to_csv(os.path.join(OUTPUT_DIR, \"training_Dataset-1-CodeCMR.csv\"))\n",
    "df_test.to_csv(os.path.join(OUTPUT_DIR, \"testing_Dataset-1-CodeCMR.csv\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3354834d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the \"selected functions\" to a JSON.\n",
    "# This is useful to limit the IDA analysis to some functions only.\n",
    "df_list = [df_training, df_validation, df_test]\n",
    "split_list = [\"training\", \"validation\", \"testing\"]\n",
    "\n",
    "for split, df_t in zip(split_list, df_list):\n",
    "\n",
    "    fset = set([tuple(x) for x in df_t[['idb_path', 'fva']].values])\n",
    "    print(\"{}: {} functions\".format(split, len(fset)))\n",
    "\n",
    "    selected_functions = defaultdict(list)\n",
    "    for t in fset:\n",
    "        selected_functions[t[0]].append(int(t[1], 16))\n",
    "        \n",
    "    # Test\n",
    "    assert(sum([len(v) for v in selected_functions.values()]) == len(fset))\n",
    "\n",
    "    # Save to file\n",
    "    with open(os.path.join(OUTPUT_DIR, \"features\", split, \"selected_{}_Dataset-1-CodeCMR.json\".format(split)), \"w\") as f_out:\n",
    "        json.dump(selected_functions, f_out)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}