From 0320c9a0b84175232cae885d57e23739526f45fd Mon Sep 17 00:00:00 2001
From: Nicola Guerrera <guerrera.nicola@gmail.com>
Date: Wed, 29 Jun 2022 11:50:29 +0200
Subject: [PATCH] Refactored tests

---
 tests/benchmark.csv | 113 +++++---------------------------------------
 tests/benchmark.py  |  53 ++++++++++-----------
 tests/firegexapi.py |  29 ++++++++++--
 3 files changed, 60 insertions(+), 135 deletions(-)

diff --git a/tests/benchmark.csv b/tests/benchmark.csv
index 442bdf5..afc345c 100755
--- a/tests/benchmark.csv
+++ b/tests/benchmark.csv
@@ -1,102 +1,11 @@
-Number of regex;Multi-Thread MB/s;Single-Thread MB/s
-0;2968,44;1655,201
-1;582,24;72,545
-2;299,866;36,222
-3;202,91;25,245
-4;150,858;18,555
-5;122,142;14,911
-6;101,656;12,398
-7;87,515;10,822
-8;76,449;9,458
-9;68,049;8,372
-10;61,615;7,563
-11;54,817;6,745
-12;51,085;6,429
-13;46,561;5,809
-14;43,522;5,518
-15;40,514;5,032
-16;38,643;4,809
-17;36,049;4,534
-18;34,211;4,249
-19;32,303;3,938
-20;30,805;3,237
-21;29,113;3,64
-22;27,837;3,484
-23;26,555;3,157
-24;25,501;3,193
-25;24,634;3,055
-26;23,286;2,928
-27;22,656;2,807
-28;21,794;2,737
-29;21,198;2,69
-30;20,387;2,504
-31;19,808;2,465
-32;19,175;2,114
-33;18,629;2,329
-34;17,785;2,282
-35;17,526;2,193
-36;16,925;2,153
-37;16,463;2,112
-38;16,136;2,02
-39;15,518;1,967
-40;15,369;1,934
-41;14,798;1,878
-42;14,451;1,821
-43;14,21;1,806
-44;13,959;1,75
-45;13,538;1,72
-46;13,377;1,686
-47;12,994;1,665
-48;12,71;1,606
-49;12,538;1,517
-50;12,239;1,537
-51;12,01;1,438
-52;11,69;1,472
-53;11,632;1,46
-54;11,325;1,373
-55;11,074;1,408
-56;10,983;1,392
-57;10,602;1,355
-58;10,521;1,335
-59;10,209;1,307
-60;10,135;1,287
-61;9,993;1,274
-62;9,86;1,259
-63;9,82;1,245
-64;9,503;1,192
-65;9,488;1,189
-66;9,292;1,145
-67;9,21;1,151
-68;9,064;1,125
-69;8,953;1,115
-70;8,646;1,088
-71;8,635;1,095
-72;8,497;1,077
-73;8,309;1,058
-74;8,247;1,05
-75;8,163;1,035
-76;8,038;1,015
-77;7,999;0,976
-78;7,837;0,983
-79;7,818;0,97
-80;7,56;0,944
-81;7,531;0,946
-82;7,497;0,95
-83;7,427;0,944
-84;7,354;0,938
-85;7,302;0,913
-86;7,053;0,905
-87;7,107;0,892
-88;6,959;0,878
-89;6,952;0,875
-90;6,898;0,863
-91;6,808;0,85
-92;6,785;0,84
-93;6,705;0,824
-94;6,518;0,824
-95;6,568;0,816
-96;6,305;0,724
-97;6,436;0,799
-98;6,232;0,771
-99;6,293;0,781
-100;6,039;0,761
+0,567.771
+1,103.886
+2,55.872
+3,37.958
+4,28.79
+5,22.956
+6,19.34
+7,16.586
+8,14.661
+9,13.158
+10,11.862
diff --git a/tests/benchmark.py b/tests/benchmark.py
index 7076ac3..decb866 100755
--- a/tests/benchmark.py
+++ b/tests/benchmark.py
@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 import argparse, socket,secrets, base64, iperf3, csv
 from time import sleep
-from requests import Session
+from firegexapi import FiregexAPI
 from multiprocessing  import Process
 
 pref = "\033["
@@ -37,31 +37,28 @@ sep()
 puts(f"Benchmarking with {args.num_of_regexes} regexes will start on ", color=colors.cyan, end="")
 puts(f"{args.address}", color=colors.yellow)
 
-s = Session()
+firegex = FiregexAPI(args.address)
+
 #Connect to Firegex
-req = s.post(f"{args.address}api/login", json={"password":args.password})
-assert req.json()["status"] == "ok", f"Benchmark Failed: Unknown response or wrong passowrd {req.text}"
-puts(f"Sucessfully logged in ✔", color=colors.green)
+if (firegex.login(args.password)): puts(f"Sucessfully logged in ✔", color=colors.green)
+else: puts(f"Benchmark Failed: Unknown response or wrong passowrd ✗", color=colors.red); exit(1)
+
 
 if args.new_istance:
     #Create new Service
-    req = s.post(f"{args.address}api/services/add" , json={"name":args.service_name,"port":args.service_port})
-    assert req.json()["status"] == "ok", f"Benchmark Failed: Couldn't create service {req.text} ✔"
-    puts(f"Sucessfully created service {args.service_name} with public port {args.service_port} ✔", color=colors.green)
+    if (firegex.create_service(args.service_name,args.service_port)):
+     puts(f"Sucessfully created service {args.service_name} with public port {args.service_port} ✔", color=colors.green)
+     service_created = True
+    else: puts(f"Benchmark Failed: Couldn't create service ✗", color=colors.red); exit(1)
 
 #Find the Service
-req = s.get(f"{args.address}api/services")
-internal_port = service_id = None
-try:
-    for service in req.json():
-        if service["name"] == args.service_name:
-            service_id = service["id"]
-            internal_port = service["internal_port"]
-            puts(f"Sucessfully received the internal port {internal_port} ✔", color=colors.green)
-            break
-except Exception:
-    puts(f"Benchmark Failed: Coulnd't get the service internal port {req.text}", color=colors.red)
-    exit(1)
+service = firegex.get_service_details(args.service_name)
+if (service):
+    internal_port= service["internal_port"]
+    service_id = service["id"]
+    puts(f"Sucessfully received the internal port {internal_port} ✔", color=colors.green)
+else: puts(f"Benchmark Failed: Coulnd't get the service internal port ✗", color=colors.red); exit_test(1)
+
 
 #Start iperf3
 def startServer():
@@ -91,9 +88,8 @@ puts(f"Baseline without proxy: ", color=colors.blue, end='')
 print(f"{getReading(internal_port)} MB/s")
 
 #Start firewall
-req = s.get(f"{args.address}api/service/{service_id}/start")
-assert req.json()["status"] == "ok", f"Benchmark Failed: Couldn't start the service {req.text}"
-puts(f"Sucessfully started service with id {service_id} ✔", color=colors.green)
+if(firegex.start(service_id)): puts(f"Sucessfully started service with id {service_id} ✔", color=colors.green)
+else: puts(f"Benchmark Failed: Coulnd't start the service ✗", color=colors.red); exit_test(1)
 
 #Hacky solution - wait a bit for the server to start
 sleep(1)
@@ -107,9 +103,7 @@ print(f"{results[0]} MB/s")
 #Add all the regexs
 for i in range(1,args.num_of_regexes+1):
     regex = base64.b64encode(bytes(secrets.token_hex(16).encode())).decode()
-    req = s.post(f"{args.address}api/regexes/add", 
-                json={"is_blacklist":True,"is_case_sensitive":True,"service_id":service_id,"mode":"B","regex":regex})
-    assert req.json()["status"] == "ok", f"Test Failed: Couldn't add regex {req.text}"
+    if(not firegex.add_regex(service_id,regex)): puts(f"Benchmark Failed: Coulnd't add the regex ✗", color=colors.red); exit_test(1)
     puts(f"Performance with {i} regex(s): ", color=colors.red, end='')
     results.append(getReading(args.service_port))
     print(f"{results[i]} MB/s")
@@ -123,8 +117,9 @@ puts(f"Sucessfully written results to {args.output_file} ✔", color=colors.mage
 
 if args.new_istance:
     #Delete the Service 
-    req = s.get(f"{args.address}api/service/{service_id}/delete")
-    assert req.json()["status"] == "ok", f"Benchmark Failed: Couldn't delete service {req.text}"
-    puts(f"Sucessfully delete service with id {service_id} ✔", color=colors.green)
+    if(firegex.delete(service_id)):
+        puts(f"Sucessfully delete service with id {service_id} ✔", color=colors.green)
+    else:
+        puts(f"Test Failed: Couldn't delete service ✗", color=colors.red); exit(1)
 
 server.terminate()
diff --git a/tests/firegexapi.py b/tests/firegexapi.py
index 117a14e..3b1f20c 100644
--- a/tests/firegexapi.py
+++ b/tests/firegexapi.py
@@ -1,13 +1,34 @@
 from requests import Session
 
+class BearerSession():
+    def __init__(self):
+        self.s = Session()
+        self.headers = {}
+
+    def post(self, endpoint, json={}, data=""):
+        headers = self.headers
+        if data:
+            headers["Content-Type"] = "application/x-www-form-urlencoded"
+        return self.s.post(endpoint, json=json, data=data, headers=headers)
+
+    def get(self, endpoint, json={}):
+        return self.s.get(endpoint, json=json, headers=self.headers)
+    
+    def setToken(self,token):
+         self.headers = {"Authorization": f"Bearer {token}"}
+
 class FiregexAPI:
     def __init__(self,address):
-        self.s = Session()
+        self.s = BearerSession()
         self.address = address
     
     def login(self,password):
-        req = self.s.post(f"{self.address}api/login", json={"password":password})
-        return req.json()["status"] == "ok"
+        req = self.s.post(f"{self.address}api/login", data=f"username=login&password={password}")
+        try : 
+            self.s.setToken(req.json()["access_token"])
+            return True
+        except Exception:
+            return False
 
     def logout(self):
         req = self.s.get(f"{self.address}api/logout")
@@ -64,4 +85,4 @@ class FiregexAPI:
     def delete(self,service_id):
         req = self.s.get(f"{self.address}api/service/{service_id}/delete")
         return req.json()["status"] == "ok"
-    
\ No newline at end of file
+