fine_tune.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Crawl Code4rena reports from Solodit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 160,
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "\n",
    "def get_solodit(page):\n",
    "    headers = {\n",
    "        'authorization': 'Token 36dc738e703c50039f3e6f03ee696730c49c54cb', # <- replace with your own token! You can find it in the network tab of your browser\n",
    "        'user-agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/112.0.0.0 Safari/537.36',\n",
    "    }\n",
    "\n",
    "    params = {\n",
    "        'source': 'Code4rena',\n",
    "        'impact': 'HIGH,MEDIUM',\n",
    "        'finder': '',\n",
    "        'protocol': '',\n",
    "        'report_date': '',\n",
    "        'min_quality_score': '0',\n",
    "        'min_general_score': '0',\n",
    "        'tags': '',\n",
    "        'bookmarked': 'False',\n",
    "        'keyword': '',\n",
    "        'page': page,\n",
    "    }\n",
    "\n",
    "    response = requests.get('https://solodit.xyz/api/issues/rest/', params=params, headers=headers)\n",
    "    return response.json()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 161,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "124"
      ]
     },
     "execution_count": 161,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "total_pages = get_solodit(1)[\"total_pages\"]\n",
    "total_pages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 234,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/wy/h6tpyrcn4szfs0598d0y5hhh0000gn/T/ipykernel_62742/2122698606.py:3: TqdmDeprecationWarning: This function will be removed in tqdm==5.0.0\n",
      "Please use `tqdm.notebook.tqdm` instead of `tqdm.tqdm_notebook`\n",
      "  for i in tqdm.tqdm_notebook(range(1, total_pages)):\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2acbbb3f231f43dc8d9dbd99f60db3d8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/123 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import tqdm\n",
    "data = []\n",
    "for i in tqdm.tqdm_notebook(range(1, total_pages+1)):\n",
    "    try:\n",
    "        data += get_solodit(i)[\"results\"]\n",
    "    except Exception as e:\n",
    "        print(e)\n",
    "        print(f\"Error in page {i}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 236,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "with open(\"solodit.json\", \"w\") as f:\n",
    "    json.dump(data, f)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parse Markdown Data & GitHub LoC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 237,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "def parse(content):\n",
    "    loc_and_vuln = content[\"content\"].split(\"# Vulnerability details\")\n",
    "    loc = loc_and_vuln[0]\n",
    "    locs = re.findall(r\"https://github.com/(.+?)/blob/(.+?)/(.+?)#(.+)\", loc)\n",
    "    vuln = \" \".join(content[\"title\"].split(\"] \")[1:])\n",
    "    return locs, vuln"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Crawl Code From GitHub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 238,
   "metadata": {},
   "outputs": [],
   "source": [
    "import functools\n",
    "\n",
    "@functools.lru_cache(maxsize=1000)\n",
    "def fetch_github(url):\n",
    "    response = requests.get(url)\n",
    "    return response.text\n",
    "\n",
    "def crawl(repo, commit_hash, file_name, line_number):\n",
    "    line_number = line_number.replace(\"L\", \"\")\n",
    "    url = f\"https://raw.github.com/{repo}/{commit_hash}/{file_name}\"\n",
    "    content = fetch_github(url)\n",
    "    lines = content.split(\"\\n\")\n",
    "    if \"-\" in line_number:\n",
    "        start, end = line_number.split(\"-\")\n",
    "        start = max(0, int(start))\n",
    "        end = min(len(lines), int(end) + 1)\n",
    "    else:\n",
    "        line_number = int(line_number)\n",
    "        start = max(0, line_number - 1)\n",
    "        end = min(len(lines), line_number + 2)\n",
    "\n",
    "    return \"\\n\".join(lines[start:end])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 242,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/var/folders/wy/h6tpyrcn4szfs0598d0y5hhh0000gn/T/ipykernel_62742/4115993144.py:1: TqdmDeprecationWarning: This function will be removed in tqdm==5.0.0\n",
      "Please use `tqdm.notebook.tqdm` instead of `tqdm.tqdm_notebook`\n",
      "  for idx, i in tqdm.tqdm_notebook(enumerate(data)):\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "86a96b6e4bed4fe6a67f0c40833d7500",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "invalid literal for int() with base 10: '199), the cidData is written directly, without checking and handling the case that a previously added nft may not have been removed:\\r'\n",
      "Error in 11 ('OpenCoreCH/cid-c4-squad', '4558d25aa8ea92644f3e778457fd6708104e0f24', 'src/CidNFT.sol', 'L192-L199), the cidData is written directly, without checking and handling the case that a previously added nft may not have been removed:\\r')\n",
      "invalid literal for int() with base 10: '199):\\r'\n",
      "Error in 11 ('OpenCoreCH/cid-c4-squad', '4558d25aa8ea92644f3e778457fd6708104e0f24', 'src/CidNFT.sol', 'L192-L199):\\r')\n",
      "invalid literal for int() with base 10: '#127'\n",
      "Error in 199 ('pooltogether/ERC5164', '5647bd84f2a6d1a37f41394874d567e45a97bf48', 'src/ethereum-arbitrum/EthereumToArbitrumRelayer.sol', 'L118-#L127')\n",
      "invalid literal for int() with base 10: '465:#466'\n",
      "Error in 789 ('code-423n4/2022-06-illuminate', 'main', 'lender/Lender.sol', 'L465:#L466')\n",
      "invalid literal for int() with base 10: '235)'\n",
      "Error in 805 ('code-423n4/2022-06-illuminate', 'main', 'lender/Lender.sol', 'L192-L235)')\n",
      "invalid literal for int() with base 10: '534)'\n",
      "Error in 805 ('code-423n4/2022-06-illuminate', 'main', 'lender/Lender.sol', 'L486-L534)')\n",
      "invalid literal for int() with base 10: '589)'\n",
      "Error in 805 ('code-423n4/2022-06-illuminate', 'main', 'lender/Lender.sol', 'L545-L589)')\n",
      "invalid literal for int() with base 10: '193)'\n",
      "Error in 809 ('code-423n4/2022-06-illuminate', 'main', 'redeemer/Redeemer.sol', 'L193)')\n",
      "invalid literal for int() with base 10: ':~:text=function%20transferToke(,%7D'\n",
      "Error in 830 ('code-423n4/2022-06-yieldy', 'main', 'src/contracts/Staking.sol', ':~:text=function%20transferToke(,%7D')\n",
      "too many values to unpack (expected 2)\n",
      "Error in 917 ('code-423n4/2022-06-connext', '4dd6149748b635f95460d4c3924c7e3fb6716967', 'contracts/contracts/core/connext/facets/BridgeFacet.sol', 'L819](https://github.com/code-423n4/2022-06-connext/blob/4dd6149748b635f95460d4c3924c7e3fb6716967/contracts/contracts/core/connext/facets/BridgeFacet.sol#L819')\n",
      "too many values to unpack (expected 2)\n",
      "Error in 933 ('code-423n4/2022-06-notional-coop', '6f8c325f604e2576e2fe257b6b57892ca181509a', 'index-coop-notional-trade-module/contracts/protocol/modules/v1/NotionalTradeModule.sol', 'L526](https://github.com/code-423n4/2022-06-notional-coop/blob/6f8c325f604e2576e2fe257b6b57892ca181509a/index-coop-notional-trade-module/contracts/protocol/modules/v1/NotionalTradeModule.sol#L526)\\r')\n",
      "invalid literal for int() with base 10: '229:#244'\n",
      "Error in 1010 ('code-423n4/2022-05-rubicon', 'main', 'contracts/RubiconRouter.sol', 'L229:#L244')\n",
      "invalid literal for int() with base 10: '492)'\n",
      "Error in 1014 ('code-423n4/2022-05-rubicon', 'main', 'contracts/RubiconRouter.sol', 'L475-L492)')\n",
      "invalid literal for int() with base 10: '48)'\n",
      "Error in 1313 ('code-423n4/2022-03-lifinance', 'main', 'src/Libraries/LibSwap.sol', 'L29-L48)')\n",
      "invalid literal for int() with base 10: '#182'\n",
      "Error in 1461 ('code-423n4/2022-02-redacted-cartel', 'main', 'contracts/RewardDistributor.sol', 'L178-#L182')\n",
      "invalid literal for int() with base 10: '#73'\n",
      "Error in 1461 ('code-423n4/2022-02-redacted-cartel', 'main', 'contracts/RewardDistributor.sol', 'L65-#L73')\n",
      "invalid literal for int() with base 10: ':~:text=uint256%20amount%20%3D%20records,_nftId%5D.reserve%20%3D%20_reserve%3B'\n",
      "Error in 1475 ('code-423n4/2022-02-nested', '69cf51d8e4eeb8bce3025db7f4f74cc565c9fad3', 'contracts/NestedRecords.sol', ':~:text=uint256%20amount%20%3D%20records,_nftId%5D.reserve%20%3D%20_reserve%3B')\n",
      "invalid literal for int() with base 10: '28\">name and symbol</a>. It is possible it set them back to an empty string, uninitializing the contract and letting the <code>initialize(..)</code> function be called again. This way, the owner may,\n",
      "Error in 1727 ('code-423n4/2021-12-amun', 'main', 'contracts/basket/contracts/facets/ERC20/ERC20Facet.sol', 'L25-L28\">name and symbol</a>. It is possible it set them back to an empty string, uninitializing the contract and letting the <code>initialize(..)</code> function be called again. This way, the owner may, for example, hide minting additional tokens. Or, after accidentally setting name and symbol to empty strings, anyone can take control over the contract and mint any number of tokens.</p>')\n"
     ]
    }
   ],
   "source": [
    "training_data = []\n",
    "total_lens = 0\n",
    "\n",
    "for idx, i in tqdm.tqdm_notebook(enumerate(data)):\n",
    "    all_code = \"\"\n",
    "    try:\n",
    "        locs, vuln = parse(i)\n",
    "    except Exception as e:\n",
    "        print(e)\n",
    "        print(f\"Error in {idx}\")\n",
    "        continue\n",
    "    # if len(vuln) > 950:\n",
    "    #     print(f\"Error in {idx} {len(vuln)}\")\n",
    "    #     continue\n",
    "    for loc in locs:\n",
    "        repo, commit_hash, file_name, line_number = loc\n",
    "        try:\n",
    "            code = crawl(repo, commit_hash, file_name, line_number)\n",
    "            all_code += code + \"\\n\"\n",
    "            total_lens += len(vuln)\n",
    "        except Exception as e:\n",
    "            print(e)\n",
    "            print(f\"Error in {idx} {loc}\")\n",
    "    if len(all_code) > 0:\n",
    "        training_data.append({\n",
    "            \"text\": all_code,\n",
    "            \"label\": vuln.strip(),\n",
    "        })"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Remove Too Lengthy Code / Vulnerability"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 245,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1440"
      ]
     },
     "execution_count": 245,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "removal = []\n",
    "for idx, i in enumerate(training_data):\n",
    "    i[\"text\"] = i[\"text\"].replace(\"\\t\", \"\").replace(\"    \", \"\")\n",
    "    if len(i[\"text\"]) < 20:\n",
    "        removal.append(idx)\n",
    "        print(i)\n",
    "\n",
    "for i in removal[::-1]:\n",
    "    del training_data[i]\n",
    "len(training_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 246,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "with open(\"training_data.pkl\", \"wb\") as f:\n",
    "    pickle.dump(training_data, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 248,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1260"
      ]
     },
     "execution_count": 248,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "removal = []\n",
    "total_lens = []\n",
    "for i in training_data:\n",
    "    if len(i[\"label\"]  + i[\"text\"]) > 1000:\n",
    "        removal.append(i)\n",
    "        total_lens.append(len(i[\"label\"] + i[\"text\"]))\n",
    "\n",
    "for i in removal[::-1]:\n",
    "    training_data.remove(i)\n",
    "len(training_data)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convert to OpenAI JSONL Format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 251,
   "metadata": {},
   "outputs": [],
   "source": [
    "fp = open(\"fine_tune.txt\", \"w\")\n",
    "for i in training_data:\n",
    "    fp.write(json.dumps(\n",
    "        {\n",
    "            \"prompt\": i[\"text\"],\n",
    "            \"completion\": i[\"label\"],\n",
    "        }\n",
    "    ))\n",
    "    fp.write(\"\\n\")\n",
    "fp.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 252,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Upload progress: 100%|███████████████████████| 490k/490k [00:00<00:00, 522Mit/s]\n",
      "Uploaded file from ./fine_tune.txt: file-0JIZiNjjmOGfI9giyONlk42g\n",
      "Created fine-tune: ft-4HRl6BNwtdUTEMklsphePbcx\n",
      "Streaming events until fine-tuning is complete...\n",
      "\n",
      "(Ctrl-C will interrupt the stream, but not cancel the fine-tune)\n",
      "[2023-04-22 22:36:09] Created fine-tune: ft-4HRl6BNwtdUTEMklsphePbcx\n",
      "[2023-04-22 22:36:20] Fine-tune costs $15.97\n",
      "[2023-04-22 22:36:21] Fine-tune enqueued. Queue number: 0\n"
     ]
    }
   ],
   "source": [
    "!openai api fine_tunes.create -t ./fine_tune.txt -m davinci"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 217,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2023-04-22 21:32:52] Created fine-tune: ft-fgUWvKK6Jpb8ZscMxmMJ5Joi\n",
      "[2023-04-22 21:33:00] Fine-tune costs $1.20\n",
      "[2023-04-22 21:33:00] Fine-tune enqueued. Queue number: 2\n",
      "[2023-04-22 21:33:39] Fine-tune is in the queue. Queue number: 1\n",
      "[2023-04-22 21:34:32] Fine-tune is in the queue. Queue number: 0\n",
      "[2023-04-22 21:35:29] Fine-tune started\n",
      "[2023-04-22 21:37:28] Completed epoch 1/4\n",
      "[2023-04-22 21:37:59] Completed epoch 2/4\n",
      "[2023-04-22 21:38:29] Completed epoch 3/4\n",
      "[2023-04-22 21:39:00] Completed epoch 4/4\n",
      "[2023-04-22 21:39:34] Uploaded model: davinci:ft-personal-2023-04-23-04-39-34\n",
      "[2023-04-22 21:39:36] Uploaded result file: file-ufOV1uNei2tozTCcbTRztBds\n",
      "[2023-04-22 21:39:36] Fine-tune succeeded\n",
      "\n",
      "Job complete! Status: succeeded 🎉\n",
      "Try out your fine-tuned model:\n",
      "\n",
      "openai api completions.create -m davinci:ft-personal-2023-04-23-04-39-34 -p <YOUR_PROMPT>\n"
     ]
    }
   ],
   "source": [
    "!openai api fine_tunes.follow -i ft-4HRl6BNwtdUTEMklsphePbcx # <- change this to your fine-tune id (starts with ft-)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_NAME = \"davinci:ft-personal-2023-04-23-04-39-34\" # <- change this to your model name from the previous step (starts with davinci:ft-)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 229,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"\"\"\n",
    "uint256 buyoutPrice = (msg.value * 100) /\n",
    "    (100 - ((depositAmount * 100) / totalSupply));\n",
    "uint256 fractionPrice = buyoutPrice / totalSupply;\n",
    "\"\"\".replace(\"\\t\", \"\").replace(\"    \", \"\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```solidity\n",
    "uint256 buyoutPrice = (msg.value * 100) /\n",
    "    (100 - ((depositAmount * 100) / totalSupply));\n",
    "uint256 fractionPrice = buyoutPrice / totalSupply;\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 233,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Buyout logic is broken due to integer division. Attacker can buy tokens for much less than expected amount.\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import openai\n",
    "openai.api_key = os.getenv(\"OPENAI_API_KEY\")\n",
    "response = openai.Completion.create( model=MODEL_NAME, prompt=prompt, temperature=0, max_tokens=1024, top_p=1, frequency_penalty=0.5, presence_penalty=0, stop=[\"\\n\", \" User:\", \" AI:\"] )\n",
    "print(response[\"choices\"][0][\"text\"])"
   ]
  }
 ],
 "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.10.8"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}