Results/notebooks/Convert CodeCMR results.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5bfc640a",
   "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",
    "#  Convert CodeCMR results                                                   #\n",
    "#                                                                            #\n",
    "##############################################################################"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5cceee46",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "import pandas as pd\n",
    "\n",
    "from scipy.spatial.distance import cosine\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "12590039",
   "metadata": {},
   "outputs": [],
   "source": [
    "def cosine_similarity(e1, e2):\n",
    "    return 1 - cosine(e1, e2)\n",
    "\n",
    "def from_str_to_numpy_array(input_str):\n",
    "    return np.fromstring(input_str[1:-1], dtype=float, sep=', ')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8b5c9ec9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_cosine_similarity(df_input):\n",
    "    sim_list = list()\n",
    "    for idx, row in tqdm(df_input.iterrows()):\n",
    "\n",
    "        if row['embeddings_1'] is np.nan or \\\n",
    "                row['embeddings_2'] is np.nan:\n",
    "            print(\"[!] Missing value in (idx:{})\".format(idx))\n",
    "            sim_list.append(0)\n",
    "            continue\n",
    "\n",
    "        e1 = from_str_to_numpy_array(row['embeddings_1'])\n",
    "        e2 = from_str_to_numpy_array(row['embeddings_2'])\n",
    "        sim_list.append(cosine_similarity(e1, e2))\n",
    "    return sim_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a6d59a1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_embedding_similarity(df_pairs, df_emb):\n",
    "    df_pairs = df_pairs.merge(df_emb,\n",
    "                              how='left',\n",
    "                              left_on=['idb_path_1', 'fva_1'],\n",
    "                              right_on=['idb_path', 'fva'])\n",
    "    df_pairs.rename(columns={'embedding': 'embeddings_1'}, inplace=True)\n",
    "    \n",
    "    df_pairs = df_pairs.merge(df_emb,\n",
    "                              how='left',\n",
    "                              left_on=['idb_path_2', 'fva_2'],\n",
    "                              right_on=['idb_path', 'fva'])\n",
    "    df_pairs.rename(columns={'embedding': 'embeddings_2'}, inplace=True)\n",
    "    \n",
    "    df_pairs['sim'] = compute_cosine_similarity(df_pairs)\n",
    "    df_pairs = df_pairs[['idb_path_1','fva_1','idb_path_2','fva_2','sim']]\n",
    "    return df_pairs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7654599f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "200000it [00:24, 8021.55it/s]\n",
      "200000it [00:24, 8028.75it/s]\n",
      "800it [00:00, 7859.95it/s]\n",
      "80000it [00:09, 8008.25it/s]\n"
     ]
    }
   ],
   "source": [
    "DB1_PATH = \"../../DBs/Dataset-1-CodeCMR/pairs/testing/\"\n",
    "df_emb = pd.read_csv(\"../data/raw_results/CodeCMR/Dataset-1-CodeCMR.csv\")\n",
    "\n",
    "df_pos = pd.read_csv(os.path.join(DB1_PATH, \"pos_testing_Dataset-1-CodeCMR.csv\"), index_col=0)\n",
    "df_neg = pd.read_csv(os.path.join(DB1_PATH, \"neg_testing_Dataset-1-CodeCMR.csv\"), index_col=0)\n",
    "df_pos_rank = pd.read_csv(os.path.join(DB1_PATH, \"pos_rank_testing_Dataset-1-CodeCMR.csv\"), index_col=0)\n",
    "df_neg_rank = pd.read_csv(os.path.join(DB1_PATH, \"neg_rank_testing_Dataset-1-CodeCMR.csv\"), index_col=0)\n",
    "\n",
    "df_pos = compute_embedding_similarity(df_pos, df_emb)\n",
    "df_neg = compute_embedding_similarity(df_neg, df_emb)\n",
    "df_pos_rank = compute_embedding_similarity(df_pos_rank, df_emb)\n",
    "df_neg_rank = compute_embedding_similarity(df_neg_rank, df_emb)\n",
    "\n",
    "df_pos.to_csv(\"../data/Dataset-1-CodeCMR/pos_testing_Dataset-1-CodeCMR.csv\", index=False)\n",
    "df_neg.to_csv(\"../data/Dataset-1-CodeCMR/neg_testing_Dataset-1-CodeCMR.csv\", index=False)\n",
    "df_pos_rank.to_csv(\"../data/Dataset-1-CodeCMR/pos_rank_testing_Dataset-1-CodeCMR.csv\", index=False)\n",
    "df_neg_rank.to_csv(\"../data/Dataset-1-CodeCMR/neg_rank_testing_Dataset-1-CodeCMR.csv\", index=False)"
   ]
  }
 ],
 "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}