CaseStudy.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Preparing the environment\n",
    "Loading models and tokenizers for CLAP and calculating SHA-256 hash of CLAP model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.multiprocessing\n",
    "import torch\n",
    "import json\n",
    "from transformers import AutoModel, AutoTokenizer\n",
    "\n",
    "device = torch.device(\"cuda\")\n",
    "\n",
    "asm_tokenizer       = AutoTokenizer.from_pretrained(\"hustcw/clap-asm\", trust_remote_code=True)\n",
    "text_tokenizer      = AutoTokenizer.from_pretrained(\"hustcw/clap-text\", trust_remote_code=True)\n",
    "asm_encoder         = AutoModel.from_pretrained(\"hustcw/clap-asm\", trust_remote_code=True).to(device)\n",
    "text_encoder        = AutoModel.from_pretrained(\"hustcw/clap-text\", trust_remote_code=True).to(device)\n",
    "\n",
    "bubble_output       = \"./CaseStudy/bubblesort.json\"\n",
    "malware_output      = \"./CaseStudy/malware.json\"\n",
    "sha3              = \"./CaseStudy/sha3.json\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fine-grained sorting algorithm classification (Zero-Shot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "bubblesort zeroshot:\n",
      "Probability: 17.954%, Text: This is a function related to bubble sort \n",
      "Probability: 6.919%, Text: This is a function related to selection sort\n",
      "Probability: 11.567%, Text: This is a function related to insertion sort\n",
      "Probability: 5.261%, Text: This is a function related to merge sort\n",
      "Probability: 9.474%, Text: This is a function related to quick sort\n",
      "Probability: 12.454%, Text: This is a function related to radix sort\n",
      "Probability: 12.879%, Text: This is a function related to shell sort\n",
      "Probability: 9.756%, Text: This is a function related to counting sort\n",
      "Probability: 9.351%, Text: This is a function related to bucket sort\n",
      "Probability: 4.385%, Text: This is a function related to heap sort\n"
     ]
    }
   ],
   "source": [
    "\n",
    "with open(bubble_output) as fp:\n",
    "    asm = json.load(fp)\n",
    "\n",
    "prompts = [\n",
    "    \"This is a function related to bubble sort \",\n",
    "    \"This is a function related to selection sort\",\n",
    "    \"This is a function related to insertion sort\",\n",
    "    \"This is a function related to merge sort\",\n",
    "    \"This is a function related to quick sort\",\n",
    "    \"This is a function related to radix sort\",\n",
    "    \"This is a function related to shell sort\",\n",
    "    \"This is a function related to counting sort\",\n",
    "    \"This is a function related to bucket sort\",\n",
    "    \"This is a function related to heap sort\",\n",
    "]\n",
    "\n",
    "with torch.no_grad():\n",
    "    asm_input = asm_tokenizer([asm], padding=True, pad_to_multiple_of=8, return_tensors=\"pt\", verbose=False)\n",
    "    asm_input = asm_input.to(device)\n",
    "    asm_embedding = asm_encoder(**asm_input)\n",
    "\n",
    "with torch.no_grad():\n",
    "    text_input = text_tokenizer(prompts, padding=True, truncation=True, return_tensors='pt')\n",
    "    text_input = text_input.to(device)\n",
    "    text_embeddings = text_encoder(**text_input)\n",
    "\n",
    "logits = torch.einsum(\"nc,ck->nk\", [asm_embedding, text_embeddings.T])\n",
    "_, preds = torch.max(logits, dim=1)\n",
    "preds = torch.softmax(logits / 0.07, dim=1).squeeze(0).tolist()\n",
    "\n",
    "print(\"bubblesort zeroshot:\")\n",
    "for i in range(len(prompts)):\n",
    "    print(f\"Probability: {round(preds[i]*100, 3)}%, Text: {prompts[i]}\")\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fine-grained malware functionality classification (Zero-Shot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "malware zeroshot:\n",
      "Probability: 75.98%, Text: This is a function related to screen shot\n",
      "Probability: 7.844%, Text: This is a function related to auto start\n",
      "Probability: 1.515%, Text: This is a function related to backdoor\n",
      "Probability: 1.616%, Text: This is a function related to download\n",
      "Probability: 2.431%, Text: This is a function related to upload\n",
      "Probability: 3.327%, Text: This is a function related to rootkit\n",
      "Probability: 1.482%, Text: This is a function related to anti detect\n",
      "Probability: 3.228%, Text: This is a function related to anti debug\n",
      "Probability: 0.919%, Text: This is a function related to passwords brute force\n",
      "Probability: 1.658%, Text: This is a function related to file hijack\n"
     ]
    }
   ],
   "source": [
    "with open(malware_output) as fp:\n",
    "    asm = json.load(fp)\n",
    "\n",
    "prompts = [\n",
    "    \"This is a function related to screen shot\",\n",
    "    \"This is a function related to auto start\",\n",
    "    \"This is a function related to backdoor\",\n",
    "    \"This is a function related to download\",\n",
    "    \"This is a function related to upload\",\n",
    "    \"This is a function related to rootkit\",\n",
    "    \"This is a function related to anti detect\",\n",
    "    \"This is a function related to anti debug\",\n",
    "    \"This is a function related to passwords brute force\",\n",
    "    \"This is a function related to file hijack\",\n",
    "]\n",
    "\n",
    "with torch.no_grad():\n",
    "    asm_input = asm_tokenizer([asm], padding=True, pad_to_multiple_of=8, return_tensors=\"pt\", verbose=False)\n",
    "    asm_input = asm_input.to(device)\n",
    "    asm_embedding = asm_encoder(**asm_input)\n",
    "\n",
    "with torch.no_grad():\n",
    "    text_input = text_tokenizer(prompts, padding=True, truncation=True, return_tensors='pt')\n",
    "    text_input = text_input.to(device)\n",
    "    text_embeddings = text_encoder(**text_input)\n",
    "\n",
    "logits = torch.einsum(\"nc,ck->nk\", [asm_embedding, text_embeddings.T])\n",
    "_, preds = torch.max(logits, dim=1)\n",
    "preds = torch.softmax(logits / 0.07, dim=1).squeeze(0).tolist()\n",
    "\n",
    "print(\"malware zeroshot:\")\n",
    "for i in range(len(prompts)):\n",
    "    print(f\"Probability: {round(preds[i]*100, 3)}%, Text: {prompts[i]}\")\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fine-grained crypto algorithm classification (Zero-Shot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sha3 zeroshot:\n",
      "Probability: 62.579%, Text: This is a function related to sha3\n",
      "Probability: 1.63%, Text: This is a function related to des\n",
      "Probability: 3.479%, Text: This is a function related to bubble sort\n",
      "Probability: 24.634%, Text: This is a function related to md5\n",
      "Probability: 5.705%, Text: This is a function related to rsa\n",
      "Probability: 1.974%, Text: This is a function related to sm4\n"
     ]
    }
   ],
   "source": [
    "with open(sha3) as fp:\n",
    "    asm = json.load(fp)\n",
    "\n",
    "prompts = [\n",
    "    \"This is a function related to sha3\",\n",
    "    \"This is a function related to des\",\n",
    "    \"This is a function related to bubble sort\",\n",
    "    \"This is a function related to md5\",\n",
    "    \"This is a function related to rsa\",\n",
    "    \"This is a function related to sm4\"\n",
    "]\n",
    "\n",
    "with torch.no_grad():\n",
    "    asm_input = asm_tokenizer([asm], padding=True, pad_to_multiple_of=8, return_tensors=\"pt\", verbose=False)\n",
    "    asm_input = asm_input.to(device)\n",
    "    asm_embedding = asm_encoder(**asm_input)\n",
    "\n",
    "with torch.no_grad():\n",
    "    text_input = text_tokenizer(prompts, padding=True, truncation=True, return_tensors='pt')\n",
    "    text_input = text_input.to(device)\n",
    "    text_embeddings = text_encoder(**text_input)\n",
    "\n",
    "logits = torch.einsum(\"nc,ck->nk\", [asm_embedding, text_embeddings.T])\n",
    "_, preds = torch.max(logits, dim=1)\n",
    "preds = torch.softmax(logits / 0.07, dim=1).squeeze(0).tolist()\n",
    "\n",
    "print(\"sha3 zeroshot:\")\n",
    "for i in range(len(prompts)):\n",
    "    print(f\"Probability: {round(preds[i]*100, 3)}%, Text: {prompts[i]}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llasm",
   "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.11.3"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}