notebooks/4.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Text Generate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "### To use gpt2, python version 3.7 or below is required.\n",
    "### To install, use : pip install tensorflow==1.15\n",
    "### Don't forget to switch kernel before running"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:\n",
      "The TensorFlow contrib module will not be included in TensorFlow 2.0.\n",
      "For more information, please see:\n",
      "  * https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md\n",
      "  * https://github.com/tensorflow/addons\n",
      "  * https://github.com/tensorflow/io (for I/O related ops)\n",
      "If you depend on functionality not listed there, please file an issue.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import gpt_2_simple as gpt2\n",
    "import os\n",
    "import requests\n",
    "import pandas as pd\n",
    "from tqdm.notebook import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading 124M model...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Fetching checkpoint: 1.05Mit [00:00, 445Mit/s]                                                      \n",
      "Fetching encoder.json: 1.05Mit [00:00, 2.50Mit/s]                                                   \n",
      "Fetching hparams.json: 1.05Mit [00:00, 250Mit/s]                                                    \n",
      "Fetching model.ckpt.data-00000-of-00001: 498Mit [01:44, 4.76Mit/s]                                  \n",
      "Fetching model.ckpt.index: 1.05Mit [00:00, 339Mit/s]                                                \n",
      "Fetching model.ckpt.meta: 1.05Mit [00:00, 3.89Mit/s]                                                \n",
      "Fetching vocab.bpe: 1.05Mit [00:00, 3.11Mit/s]                                                      \n"
     ]
    }
   ],
   "source": [
    "model_name = \"124M\"\n",
    "if not os.path.isdir(os.path.join(\"Downloads\",\"models\", model_name)):\n",
    "    print(f\"Downloading {model_name} model...\")\n",
    "    gpt2.download_gpt2(model_dir='Downloads/models', model_name=model_name)   # model is saved into current directory under /models/124M/\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Input Email TTR'ed Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def grt_txt(file_name,model_name,savepath):\n",
    "    sess = gpt2.start_tf_sess()\n",
    "    gpt2.finetune(sess,\n",
    "                  file_name,\n",
    "                  model_name=model_name,\n",
    "                  model_dir='Downloads/models',\n",
    "                  steps=20)   # steps is max number of training steps\n",
    "#     Take too much time for large number of steps... Can use 1 to 1000.\n",
    "\n",
    "    for i in tqdm(range(200)):\n",
    "        single_text = gpt2.generate(sess, return_as_list=True)\n",
    "        with open(f'{savepath}/{i}.txt','w') as file:\n",
    "            file.write(single_text[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv('../data/additional-features-v2/new/additional_features_TTR.tsv',sep='\\t')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# text = grt_txt(\"../tests/test.txt\", '124M')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# len(text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## This part is to get the 200 inital emails of each type of attack\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Index(['Unnamed: 0', 'Unnamed: 0.1', 'date', 'location', 'now frequency',\n",
       "       'time', 'urgent frequency', 'wrong cap ratio', 'attacker estimated age',\n",
       "       'attacker relationship', 'credential_phishing', 'female count',\n",
       "       'has offer', 'high total cholesterol rate', 'hypertension rate', 'ip',\n",
       "       'male count', 'malware', 'obesity rate', 'reconnaissance',\n",
       "       'sentiment analysis', 'social_engineering', 'stock_close', 'stock_open',\n",
       "       'stock_volume', 'title', 'unknown count', 'wrongspellnumber', 'ip_risk',\n",
       "       'ip_score', 'TTR'ed Text'],\n",
       "      dtype='object')"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "types = df[['reconnaissance','social_engineering','malware','credential_phishing']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>reconnaissance</th>\n",
       "      <th>social_engineering</th>\n",
       "      <th>malware</th>\n",
       "      <th>credential_phishing</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3987</th>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3988</th>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3989</th>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3990</th>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3991</th>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>3992 rows × 4 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "      reconnaissance  social_engineering  malware  credential_phishing\n",
       "0                1.0                 1.0      0.0                  1.0\n",
       "1                1.0                 1.0      0.0                  1.0\n",
       "2                0.0                 1.0      0.0                  0.0\n",
       "3                0.0                 1.0      0.0                  0.0\n",
       "4                0.0                 1.0      0.0                  1.0\n",
       "...              ...                 ...      ...                  ...\n",
       "3987             0.0                 1.0      0.0                  0.0\n",
       "3988             0.0                 1.0      0.0                  1.0\n",
       "3989             0.0                 1.0      0.0                  1.0\n",
       "3990             0.0                 1.0      0.0                  0.0\n",
       "3991             0.0                 0.0      0.0                  0.0\n",
       "\n",
       "[3992 rows x 4 columns]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "value1 = types[types.sum(axis=1)==1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Reconnaissance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "value2_r = types[(types['reconnaissance']==1) & (types.sum(axis=1)==2)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "R1_index = value1[value1['reconnaissance']==1].index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "R_list = []\n",
    "for i in R1_index:\n",
    "    R_list.append(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "R2_index = value2_r.index\n",
    "temp = 0\n",
    "while len(R_list)< 200:\n",
    "    R_list.append(R2_index[temp])\n",
    "    temp += 1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "200"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(R_list)\n",
    "# R_list is the index of emails with attack reconnaissance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "R_txt = df[\"TTR'ed Text\"].loc[R_list]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Social Engineering\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "S1_index = value1[value1['social_engineering']==1].index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "S_list = []\n",
    "for i in range(200):\n",
    "    S_list.append(S1_index[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "200"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(S_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "S_txt = df[\"TTR'ed Text\"].loc[S_list]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Malware"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "M1_index = value1[value1['malware']==1].index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "value2_m = types[(types['malware']==1) & (types.sum(axis=1)==2)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "value3_m = types[(types['malware']==1) & (types.sum(axis=1)==3)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "M_list = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in M1_index:\n",
    "    M_list.append(i)\n",
    "for j in value2_m.index:\n",
    "    M_list.append(j)\n",
    "for k in value3_m.index:\n",
    "    if len(M_list) < 200:\n",
    "        M_list.append(k)\n",
    "    else:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "200"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(M_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "M_txt = df[\"TTR'ed Text\"].loc[M_list]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Credential Phishing\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "C1_index = value1[value1['credential_phishing']==1].index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "C_list = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in C1_index:\n",
    "    if len(C_list)<200:\n",
    "        C_list.append(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "200"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(C_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "C_txt = df[\"TTR'ed Text\"].loc[C_list]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create Emails txt with different Attack type\n",
    "#### Delete the existing txt files before re-doing these\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "attack=['reconnaissance','social_engineering','malware','credential_phishing']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reconnaissance\n",
    "# with open(f'../data/additional-features-v2/new/attack_Text/reconnaissance.txt','a+') as f:\n",
    "#     for i in R_txt:\n",
    "#         f.write(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Social Engineering\n",
    "# with open(f'../data/additional-features-v2/new/attack_Text/social_engineering.txt','a+') as f:\n",
    "#     for i in S_txt:\n",
    "#         f.write(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Malware\n",
    "# with open(f'../data/additional-features-v2/new/attack_Text/malware.txt','a+') as f:\n",
    "#     for i in M_txt:\n",
    "#         f.write(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Credential_phishing\n",
    "# with open(f'../data/additional-features-v2/new/attack_Text/Credential_phishing.txt','a+') as f:\n",
    "#     for i in C_txt:\n",
    "#         f.write(i)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Text generating"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /Users/xiaoyuwang/.local/share/virtualenvs/DSCI-550-Assignment-2-f1o6Jtui/lib/python3.7/site-packages/gpt_2_simple/src/sample.py:17: where (from tensorflow.python.ops.array_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use tf.where in 2.0, which has the same broadcast rule as np.where\n",
      "Loading checkpoint checkpoint/run1/model-5\n",
      "INFO:tensorflow:Restoring parameters from checkpoint/run1/model-5\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/1 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading dataset...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:00<00:00,  1.25it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dataset has 116204 tokens\n",
      "Training...\n",
      "[6 | 45.78] loss=3.37 avg=3.37\n",
      "Saving checkpoint/run1/model-6\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "76918846c7ba4bcb946e5923dd04ec28",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/200 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "grt_txt(\"../data/additional-features-v2/new/attack_Text/reconnaissance.txt\", '124M','../data/additional-features-v2/new/Generated_text/Reconnaissance/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grt_txt(\"../data/additional-features-v2/new/attack_Text/social_engineering.txt\", '124M','../data/additional-features-v2/new/Generated_text/Social_engineering/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grt_txt(\"../data/additional-features-v2/new/attack_Text/malware.txt\", '124M','../data/additional-features-v2/new/Generated_text/Malware/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grt_txt(\"../data/additional-features-v2/new/attack_Text/Credential_phishing.txt\", '124M','../data/additional-features-v2/new/Generated_text/Credential_phishing/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "DSCI-550-Assignment-2",
   "language": "python",
   "name": "dsci-550-assignment-2"
  },
  "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}