DBs/Dataset-1-CodeCMR/Dataset-1-CodeCMR sanity check.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5894b9e2",
   "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 sanity check                                            #\n",
    "#                                                                            #\n",
    "##############################################################################"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "54dea881",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Requirements\n",
    "# pandas==1.4.2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "bc745ada",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "import pandas as pd\n",
    "import pickle\n",
    "import unittest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4cb6375e",
   "metadata": {},
   "outputs": [],
   "source": [
    "N_SELECTED_FUNCS_TRAINING = 181907\n",
    "N_SELECTED_BINARIES_TRAINING = 2114\n",
    "\n",
    "N_SELECTED_FUNCS_VALIDATION = 8233\n",
    "N_SELECTED_BINARIES_VALIDATION = 562\n",
    "\n",
    "N_SELECTED_FUNCS_TESTING = 286950\n",
    "N_SELECTED_BINARIES_TESTING = 600"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "fb066bcc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset-1-CodeCMR example.ipynb      \u001b[34mpairs\u001b[m\u001b[m\r\n",
      "Dataset-1-CodeCMR plots.ipynb        testing_Dataset-1-CodeCMR.csv\r\n",
      "Dataset-1-CodeCMR sanity check.ipynb training_Dataset-1-CodeCMR.csv\r\n",
      "\u001b[34mfeatures\u001b[m\u001b[m                             validation_Dataset-1-CodeCMR.csv\r\n"
     ]
    }
   ],
   "source": [
    "!ls ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "131eccb6",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "test_binary_files (__main__.TestDataset1) ... ok\n",
      "test_pairs_dataset_testing (__main__.TestDataset1) ... ok\n",
      "test_pairs_dataset_validation (__main__.TestDataset1) ... ok\n",
      "test_pickle_files (__main__.TestDataset1) ... ok\n",
      "test_selected (__main__.TestDataset1) ... ok\n",
      "test_testing_dataset (__main__.TestDataset1) ... ok\n",
      "test_training_dataset (__main__.TestDataset1) ... ok\n",
      "test_validation_dataset (__main__.TestDataset1) ... ok\n",
      "\n",
      "----------------------------------------------------------------------\n",
      "Ran 8 tests in 15.825s\n",
      "\n",
      "OK\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<unittest.main.TestProgram at 0x117ab95b0>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class TestDataset1(unittest.TestCase):\n",
    "    \n",
    "    def setUp(self):\n",
    "        training_dataset = pd.read_csv(\"training_Dataset-1-CodeCMR.csv\", index_col=0)\n",
    "        self.fva_set_training = set([tuple(x) for x in training_dataset[['idb_path', 'fva']].values])\n",
    "        self.idb_set_training = set(training_dataset['idb_path'].values)\n",
    "        self.pickle_set_training = set(training_dataset['pickle_path'].values)\n",
    "        \n",
    "        validation_dataset = pd.read_csv(\"validation_Dataset-1-CodeCMR.csv\", index_col=0)\n",
    "        self.fva_set_validation = set([tuple(x) for x in validation_dataset[['idb_path', 'fva']].values])\n",
    "        self.idb_set_validation = set(validation_dataset['idb_path'].values)\n",
    "        self.pickle_set_validation = set(validation_dataset['pickle_path'].values)\n",
    "\n",
    "        testing_dataset = pd.read_csv(\"testing_Dataset-1-CodeCMR.csv\", index_col=0)\n",
    "        self.fva_set_testing = set([tuple(x) for x in testing_dataset[['idb_path', 'fva']].values])\n",
    "        self.idb_set_testing = set(testing_dataset['idb_path'].values)\n",
    "        self.pickle_set_testing = set(testing_dataset['pickle_path'].values)\n",
    "    \n",
    "    def test_training_dataset(self):\n",
    "        training_dataset = pd.read_csv(\"training_Dataset-1-CodeCMR.csv\", index_col=0)\n",
    "\n",
    "        # Test null values\n",
    "        self.assertFalse(training_dataset.isnull().values.any())\n",
    "\n",
    "        # Test number of unique functions\n",
    "        self.assertEqual(training_dataset.shape[0], N_SELECTED_FUNCS_TRAINING)\n",
    "        fva_set_training = set([tuple(x) for x in training_dataset[['idb_path', 'fva']].values])\n",
    "        self.assertEqual(len(fva_set_training), N_SELECTED_FUNCS_TRAINING)\n",
    "\n",
    "        # Test the number of unique binaries\n",
    "        idb_set_training = set(training_dataset['idb_path'].values)\n",
    "        self.assertEqual(len(idb_set_training), N_SELECTED_BINARIES_TRAINING)\n",
    "    \n",
    "    def test_validation_dataset(self):\n",
    "        validation_dataset = pd.read_csv(\"validation_Dataset-1-CodeCMR.csv\", index_col=0)\n",
    "        \n",
    "        # Test null values\n",
    "        self.assertFalse(validation_dataset.isnull().values.any())\n",
    "        \n",
    "        # Test number of unique functions\n",
    "        self.assertEqual(validation_dataset.shape[0], N_SELECTED_FUNCS_VALIDATION)\n",
    "        fva_set_validation = set([tuple(x) for x in validation_dataset[['idb_path', 'fva']].values])\n",
    "        self.assertEqual(len(fva_set_validation), N_SELECTED_FUNCS_VALIDATION)\n",
    "        \n",
    "        # Test the number of unique binaries\n",
    "        idb_set_validation = set(validation_dataset['idb_path'].values)\n",
    "        self.assertEqual(len(idb_set_validation), N_SELECTED_BINARIES_VALIDATION)\n",
    "    \n",
    "    def test_testing_dataset(self):\n",
    "        testing_dataset = pd.read_csv(\"testing_Dataset-1-CodeCMR.csv\", index_col=0)\n",
    "\n",
    "        # Test null values\n",
    "        self.assertFalse(testing_dataset.isnull().values.any())\n",
    "\n",
    "        # Test number of unique functions\n",
    "        self.assertEqual(testing_dataset.shape[0], N_SELECTED_FUNCS_TESTING)\n",
    "        fva_set_testing = set([tuple(x) for x in testing_dataset[['idb_path', 'fva']].values])\n",
    "        self.assertEqual(len(fva_set_testing), N_SELECTED_FUNCS_TESTING)\n",
    "\n",
    "        # Test the number of unique binaries\n",
    "        idb_set_testing = set(testing_dataset['idb_path'].values)\n",
    "        self.assertEqual(len(idb_set_testing), N_SELECTED_BINARIES_TESTING)\n",
    "\n",
    "    def test_pairs_dataset_testing(self):\n",
    "        for pair_name in os.listdir(\"pairs/testing/\"):\n",
    "            if not pair_name.endswith(\".csv\"):\n",
    "                continue\n",
    "            pair_path = os.path.join(\"pairs/testing/\", pair_name)\n",
    "\n",
    "            pairs_dataset = pd.read_csv(pair_path, index_col=0)\n",
    "\n",
    "            # Test null values\n",
    "            self.assertFalse(pairs_dataset.isnull().values.any())\n",
    "\n",
    "            # Test pairs size            \n",
    "            dt = dict(pairs_dataset.groupby(\"db_type\").count()['idb_path_1'].items())\n",
    "            if pair_name.startswith(\"pos_rank\"):\n",
    "                self.assertDictEqual(dt, {'XA': 200, 'XC': 200, 'XC+XB': 200, 'XM': 200})\n",
    "            elif pair_name.startswith(\"neg_rank\"):\n",
    "                self.assertDictEqual(dt, {'XA': 20000, 'XC': 20000, 'XC+XB': 20000, 'XM': 20000})\n",
    "            else:\n",
    "                self.assertDictEqual(dt, {'XA': 50000, 'XC': 50000, 'XC+XB': 50000, 'XM': 50000})\n",
    "\n",
    "            # Test overlapping functions\n",
    "            p_fva_set = set([tuple(x) for x in pairs_dataset[['idb_path_1', 'fva_1']].values])\n",
    "            p_fva_set |= set([tuple(x) for x in pairs_dataset[['idb_path_2', 'fva_2']].values])\n",
    "            self.assertEqual(p_fva_set & self.fva_set_testing, p_fva_set)\n",
    "    \n",
    "    def test_pairs_dataset_validation(self):\n",
    "        for pair_name in os.listdir(\"pairs/validation/\"):\n",
    "            if not pair_name.endswith(\".csv\"):\n",
    "                continue\n",
    "            pair_path = os.path.join(\"pairs/validation/\", pair_name)\n",
    "\n",
    "            pairs_dataset = pd.read_csv(pair_path, index_col=0)\n",
    "\n",
    "            # Test null values\n",
    "            self.assertFalse(pairs_dataset.isnull().values.any())\n",
    "\n",
    "            # Test pairs size            \n",
    "            dt = dict(pairs_dataset.groupby(\"db_type\").count()['idb_path_1'].items())\n",
    "            self.assertDictEqual(dt, {'XA': 10000, 'XC': 10000, 'XC+XB': 10000, 'XM': 10000})\n",
    "\n",
    "            # Test overlapping functions\n",
    "            p_fva_set = set([tuple(x) for x in pairs_dataset[['idb_path_1', 'fva_1']].values])\n",
    "            p_fva_set |= set([tuple(x) for x in pairs_dataset[['idb_path_2', 'fva_2']].values])\n",
    "            self.assertEqual(p_fva_set & self.fva_set_validation, p_fva_set)\n",
    "    \n",
    "    def test_selected(self):\n",
    "        with open(\"features/training/selected_training_Dataset-1-CodeCMR.json\") as f_in:\n",
    "            selected = json.load(f_in)\n",
    "            \n",
    "            # Test overlapping functions\n",
    "            entries_s = set([(k, hex(v)) for k in selected for v in selected[k]])\n",
    "            self.assertEqual(entries_s & self.fva_set_training, entries_s | self.fva_set_training)\n",
    "        \n",
    "        with open(\"features/validation/selected_validation_Dataset-1-CodeCMR.json\") as f_in:\n",
    "            selected = json.load(f_in)\n",
    "            \n",
    "            # Test overlapping functions\n",
    "            entries_s = set([(k, hex(v)) for k in selected for v in selected[k]])\n",
    "            self.assertEqual(entries_s & self.fva_set_validation, entries_s | self.fva_set_validation)\n",
    "        \n",
    "        with open(\"features/testing/selected_testing_Dataset-1-CodeCMR.json\") as f_in:\n",
    "            selected = json.load(f_in)\n",
    "            \n",
    "            # Test overlapping functions\n",
    "            entries_s = set([(k, hex(v)) for k in selected for v in selected[k]])\n",
    "            self.assertEqual(entries_s & self.fva_set_testing, entries_s | self.fva_set_testing)\n",
    "            \n",
    "    def test_binary_files(self):\n",
    "        a_list = [self.idb_set_training, self.idb_set_validation, self.idb_set_testing]\n",
    "        b_list = [N_SELECTED_BINARIES_TRAINING, N_SELECTED_BINARIES_VALIDATION, N_SELECTED_BINARIES_TESTING]\n",
    "        for bin_list, n_binaries in zip (a_list, b_list):\n",
    "            binary_counter = 0\n",
    "            for path in bin_list:\n",
    "                npath = path.replace(\"IDBs/\", \"../../Binaries/\")\n",
    "                npath = npath.replace(\".i64\", \"\")\n",
    "                if os.path.isfile(npath):\n",
    "                    binary_counter += 1\n",
    "            self.assertEqual(binary_counter, n_binaries)\n",
    "            \n",
    "    def test_pickle_files(self):\n",
    "        a_list = [self.pickle_set_training, self.pickle_set_validation, self.pickle_set_testing]\n",
    "        b_list = [N_SELECTED_BINARIES_TRAINING, N_SELECTED_BINARIES_VALIDATION, N_SELECTED_BINARIES_TESTING]\n",
    "        \n",
    "        for pickle_list, n_binaries in zip (a_list, b_list):\n",
    "            pickle_counter = 0\n",
    "            for pickle_path in pickle_list:\n",
    "                if os.path.isfile(pickle_path):\n",
    "                    pickle_counter += 1\n",
    "            self.assertEqual(pickle_counter, n_binaries)\n",
    "        \n",
    "unittest.main(argv=[''], verbosity=2, exit=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
}