diff --git a/README.md b/README.md index c1a5dbd..578e5b0 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,20 @@ -# walle +# Titli Artificial Intelligence based Intrusion Detection Systems + +![PyPI - Python Version](https://img.shields.io/pypi/pyversions/titli) +![PyPI - Version](https://img.shields.io/pypi/v/titli) + + +### Installation +``` +pip install titli +``` + +### Usage +- Step 1: Copy the ```examples/train_ids.py``` and ```examples/test_ids.py``` file from the repo to your local machine. +- Step 2: Run both the files to train and test the Kitsune IDS respectively. + +### Todo (Developer Tasks) +- [ ] Check if RMSE is used for loss or just the difference. +- [ ] Put Kitsune code into the base IDS format. +- [ ] Write code to evaluate the model and calculate all the metrics. diff --git a/backup/kitsune_plot.png b/backup/kitsune_plot.png deleted file mode 100644 index 7e8eaa8..0000000 Binary files a/backup/kitsune_plot.png and /dev/null differ diff --git a/backup/raids.zip b/backup/raids.zip deleted file mode 100644 index 0fa5b0b..0000000 Binary files a/backup/raids.zip and /dev/null differ diff --git a/backup/weekday_raw_re.png b/backup/weekday_raw_re.png deleted file mode 100644 index 89847d9..0000000 Binary files a/backup/weekday_raw_re.png and /dev/null differ diff --git a/backup/test_attack.py b/examples/test_attack.py similarity index 100% rename from backup/test_attack.py rename to examples/test_attack.py diff --git a/backup/test_ids.py b/examples/test_ids.py similarity index 97% rename from backup/test_ids.py rename to examples/test_ids.py index 99ac6e3..c13d064 100644 --- a/backup/test_ids.py +++ b/examples/test_ids.py @@ -72,7 +72,7 @@ except: ax1.scatter(x_val, anomaly_scores, s=1, alpha=1.0, c='#FF8C00') -ax1.axhline(y=threshold, color='r', linestyle='-') +ax1.axhline(y=threshold, color='b', linestyle='-') ax1.set_yscale("log") ax1.set_title("Anomaly Scores from Kitsune Execution Phase") ax1.set_ylabel("RMSE (log scaled)") diff --git a/backup/train_ids.py b/examples/train_ids.py similarity index 100% rename from backup/train_ids.py rename to examples/train_ids.py diff --git a/backup/train_test_pytorch_kitsune.py b/examples/train_test_pytorch_kitsune.py similarity index 75% rename from backup/train_test_pytorch_kitsune.py rename to examples/train_test_pytorch_kitsune.py index 15b49c2..8bc9cb1 100644 --- a/backup/train_test_pytorch_kitsune.py +++ b/examples/train_test_pytorch_kitsune.py @@ -1,7 +1,7 @@ -from walle.ids import PyTorchKitsune +from titli.ids import PyTorchKitsune model = PyTorchKitsune(num_features=100) threshold = model.train(pcap_path="../PANDA/PANDA/data/benign/weekday.pcap") print(threshold) model.infer(pcap_path="../PANDA/PANDA/data/benign/weekday.pcap") -model.plot_kitsune(threshold) \ No newline at end of file +model.plot_kitsune(threshold) diff --git a/titli/ids/pytorch_kitsune.py b/titli/ids/pytorch_kitsune.py index 622425b..ef99de0 100644 --- a/titli/ids/pytorch_kitsune.py +++ b/titli/ids/pytorch_kitsune.py @@ -143,6 +143,7 @@ def train(self, pcap_path: str) -> None: self.device = torch.device("mps" if torch.cuda.is_available() else "cpu") self.rmse = RMSELoss() self.FMgrace = np.floor(self.FMgrace_rate * len(rdpcap(pcap_path))) + self.dataset = pcap_path.split("/")[-1].split(".")[0] print("Training Feature Mapping (FM) phase") for i, packet in enumerate(PcapReader(pcap_path)): @@ -179,10 +180,15 @@ def train(self, pcap_path: str) -> None: loss.backward() optimiser.step() - # Making the decoder weight the transpose of the encoder weight + # # Making the decoder weight the transpose of the encoder weight + # for tail in self.tails: + # tail.decoder.weight.data = tail.encoder.weight.data.t() + # self.head.decoder.weight.data = self.head.encoder.weight.data.t() + + # Making the encoder weight the transpose of the decoder weight for tail in self.tails: - tail.decoder.weight.data = tail.encoder.weight.data.t() - self.head.decoder.weight.data = self.head.encoder.weight.data.t() + tail.encoder.weight.data = tail.decoder.weight.data.t() + self.head.encoder.weight.data = self.head.decoder.weight.data.t() if (i + 1) % self.print_interval == 0: print(f"Packet: {i} | Loss: {loss.data}") @@ -212,6 +218,8 @@ def get_threshold(self, pcap_path: str) -> float: # loss loss = self.rmse(x_hat, tails) + if loss.data == 0: + loss.data = torch.tensor(1e-2) self.threshold_rmse.append(loss.data) log_re = np.log(self.threshold_rmse) @@ -244,33 +252,29 @@ def infer(self, pcap_path: str) -> None: # loss loss = self.rmse(x_hat, tails) + if loss.data == 0: + loss.data = torch.tensor(1e-2) self.rmse_array.append(loss.data) if (i + 1) % self.print_interval == 0: print(f"Inferencing: {i} | Loss: {loss.data}") - def plot_kitsune(self,threshold,plot_with_time = True, meta_file = False, out_image = "kitsune_plot.png"): - cmap = plt.get_cmap('Set3') - - f, ax1 = plt.subplots(constrained_layout=True, figsize=(10, 5), dpi=200) + def plot_kitsune(self, threshold, out_image = "kitsune_plot.png"): + _ = plt.get_cmap('Set3') - if plot_with_time: - x_val = range(len(self.rmse_array)) - ax1.xaxis.set_major_locator(ticker.MultipleLocator(0.85)) - ax1.tick_params(labelrotation=90) - else: - x_val = range(len(self.rmse_array)) + f, ax1 = plt.subplots(constrained_layout=True, figsize=(10, 5), dpi=600) + x_val = np.arange(len(self.rmse_array)) - if meta_file: + try: ax1.scatter(x_val, self.rmse_array, s=1, c='#00008B') - else: + except: ax1.scatter(x_val, self.rmse_array, s=1, alpha=1.0, c='#FF8C00') - ax1.axhline(y=threshold, color='r', linestyle='-') + ax1.axhline(y=threshold, color='b', linestyle='-') ax1.set_yscale("log") - ax1.set_title("Anomaly Scores from Kitsune Execution Phase") - ax1.set_ylabel("RMSE (log scaled)") - ax1.set_xlabel("packet index") + ax1.set_title(f"Anomaly Scores of Kitsune Execution Phase: {self.dataset}") + ax1.set_ylabel("RMSE") + ax1.set_xlabel("Packet Index") f.savefig(out_image) print("plot path:", out_image)