test.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test SAFEtorch pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.function_normalizer import FunctionNormalizer\n",
    "from utils.instructions_converter import InstructionsConverter\n",
    "from utils.capstone_disassembler import disassemble\n",
    "from safetorch.safe_network import SAFE\n",
    "from safetorch.parameters import Config\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize SAFE\n",
    "config = Config()\n",
    "safe = SAFE(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load instruction converter and normalizer\n",
    "I2V_FILENAME = \"model/word2id.json\"\n",
    "converter = InstructionsConverter(I2V_FILENAME)\n",
    "normalizer = FunctionNormalizer(max_instruction=150)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load SAFE weights\n",
    "SAFE_torch_model_path = \"model/SAFEtorch.pt\"\n",
    "state_dict = torch.load(SAFE_torch_model_path)\n",
    "safe.load_state_dict(state_dict)\n",
    "safe = safe.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "asm = \"554889e5534881ec980000004889bd68ffffff4889b560ffffff488b8568ffffffc6808800000001488b8560ffffff4889c7e82bbf0000488945d048c745e80000000048c745e00000000048c745d800000000488d8570ffffff488b55d04889d64889c7e80dfcffff488d8570ffffff4889c7e8e67cfdff4889c2488b8568ffffff48895018488d8570ffffff4889c7e82d6bfdff488b8568ffffff488b40184889c7e8feaed6ff488945e8488b8568ffffff488b5818488d4580488b55d04889d64889c7e8acfbffff488b75e8488d4580b9000000004889da4889c7e8f49cfdff488d45804889c7e8d46afdff488b8568ffffff488b55e848895010488b8568ffffff488b40184885c0742b488b8568ffffff488b40100fb60084c07e19488b8568ffffff488b40100fb6003c7f7407b801000000eb05b800000000488b9568ffffff884220488d4590488b55d04889d64889c7e850fbffff488d45904889c7e8f87bfdff4889c2488b8568ffffff48895030488d45904889c7e8426afdff488b8568ffffff488b40304889c7e813aed6ff488945e0488b8568ffffff488b5830488d45a0488b55d04889d64889c7e8f5faffff488b75e0488d45a0b9000000004889da4889c7e8099cfdff488d45a04889c7e8e969fdff488b8568ffffff488b55e048895028488d45b0488b55d04889d64889c7e8e3faffff488d45b04889c7e8577bfdff4889c2488b8568ffffff48895040488d45b04889c7e8a169fdff488b8568ffffff488b40404889c7e872add6ff488945d8488b8568ffffff488b5840488d45c0488b55d04889d64889c7e888faffff488b75d8488d45c0b9000000004889da4889c7e8689bfdff488d45c04889c7e84869fdff488b8568ffffff488b55d848895038488b45d04889c7e895f9ffff89c2488b8568ffffff885048488b45d04889c7e8a3f9ffff89c2488b8568ffffff885049488b8560ffffff4889c7e8c0ba0000488945c8488b8568ffffff488d484a488b0590bb3200488b00488d5024488b0582bb3200488b30488b45c84889c7e88faeffff488b8568ffffff488d486e488b05994f3200488b00488d501a488b058b4f3200488b30488b45c84889c7e860aeffffe9d40000004889c3488d8570ffffff4889c7e88168fdff4889d8eb774889c3488d45804889c7e86d68fdff4889d8eb634889c3488d45904889c7e85968fdff4889d8eb4f4889c3488d45a04889c7e84568fdff4889d8eb3b4889c3488d45b04889c7e83168fdff4889d8eb274889c3488d45c04889c7e81d68fdff4889d8eb134889c3e858d5fbff4889d84889c7e8b93dccff4889c7e865d4fbff48837de800740c488b45e84889c7e82eacd6ff48837de000740c488b45e04889c7e81bacd6ff48837dd800740c488b45d84889c7e808acd6ffe833cffbff4881c4980000005b5dc3\"\n",
    "arch = \"x86\"\n",
    "bits = 64\n",
    "instructions = disassemble(asm, arch, bits)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 1.0335e-01, -5.4050e-02, -5.7076e-02, -2.1117e-01,  2.2357e-01,\n",
       "          2.0404e-02,  3.9656e-02, -1.7016e-01, -3.9423e-02, -1.3362e-01,\n",
       "          6.2012e-02,  1.2293e-02, -4.0920e-02, -1.1634e-02,  1.0139e-02,\n",
       "         -6.3374e-02, -1.2039e-02,  1.0519e-01,  5.5035e-02, -1.9403e-02,\n",
       "          4.3240e-03, -2.1956e-02,  3.5674e-02, -8.3240e-02,  7.2140e-02,\n",
       "         -1.1527e-01, -2.1823e-01,  3.5691e-02, -6.6146e-02, -2.0117e-01,\n",
       "          1.4557e-02,  1.3812e-01, -3.0271e-02,  1.8191e-02, -6.1304e-02,\n",
       "          4.8532e-02,  8.1854e-03,  4.3202e-01, -3.2430e-04,  1.8965e-02,\n",
       "          8.3106e-02, -1.8471e-02, -6.7362e-02, -1.1174e-01,  2.0485e-02,\n",
       "         -1.3074e-01,  1.2108e-01, -7.1991e-02, -1.9878e-01, -9.3574e-02,\n",
       "         -6.8292e-02,  3.9065e-02,  2.9129e-02,  1.2143e-02,  1.0071e-02,\n",
       "         -1.2389e-01,  1.6163e-01,  4.2207e-02,  4.6166e-02, -6.2232e-02,\n",
       "          1.0087e-01, -1.0922e-01,  8.2758e-02, -5.8365e-02,  1.0083e-01,\n",
       "         -1.6366e-02, -8.6855e-02, -6.8850e-02,  2.3933e-02, -1.8436e-02,\n",
       "          1.1112e-02, -7.8794e-02, -1.0857e-01,  4.5918e-02, -1.7800e-02,\n",
       "         -6.8207e-02, -2.9228e-01,  8.7067e-02,  9.5397e-02, -3.2472e-02,\n",
       "          1.4847e-02, -8.2678e-02,  1.1685e-02, -1.9182e-02,  1.6823e-02,\n",
       "          6.0362e-02, -6.8276e-02, -5.5158e-02, -1.4592e-01,  3.1696e-02,\n",
       "         -5.9405e-02,  1.3099e-02,  5.2503e-03, -8.0928e-02,  8.2340e-02,\n",
       "         -7.4340e-02, -1.6591e-01, -1.7529e-01,  7.4791e-02, -1.7182e-02]],\n",
       "       grad_fn=<DivBackward0>)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "converted_instructions = converter.convert_to_ids(instructions)\n",
    "instructions, length = normalizer.normalize_functions([converted_instructions])\n",
    "tensor = torch.LongTensor(instructions[0])\n",
    "function_embedding = safe(tensor, length)\n",
    "function_embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}